fixes mediaforensics/medifor#41

xkortex · Jul 22, 2020 · 7bf089c · 7bf089c
1 parent 946a2bf
commit 7bf089c
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 6 deletions.
diff --git a/python/examples/example_imgmanip.py b/python/examples/example_imgmanip.py
@@ -23,6 +23,7 @@ def run(img, out, foo, bar, tle, option=None):
             "mask": {"uri": out + "/cli/mask.png", "type": "image/png"},
             "maskOptout": {"uri": out + "/cli/omask.png", "type": "image/png"},
         },
+        "supplement": [{"uri": out + "/cli/supplement.txt", "type": "text/plain"}],
     }
 
     print(json.dumps(d))

diff --git a/python/medifor/v1/mediforclient.py b/python/medifor/v1/mediforclient.py
@@ -274,8 +274,16 @@ def img_manip(self, img, output_dir):
         req.request_id = str(uuid.uuid4())
         req.out_dir = output_dir
         logging.info(req)
+        resp = self.detect_one(req)
+        if resp.localization.mask.uri:
+            resp.localization.mask.uri = self.o_unmap(resp.localization.mask.uri)
+        if resp.localization.mask_optout.uri:
+            resp.localization.mask_optout.uri = self.o_unmap(resp.localization.mask_optout.uri)
 
-        return self.detect_one(req)
+        for msg in resp.supplement:
+            if msg.uri:
+                msg.uri = self.o_unmap(msg.uri)
+        return resp
 
     def vid_manip(self, vid, output_dir):
         """
@@ -298,7 +306,17 @@ def vid_manip(self, vid, output_dir):
         req.out_dir = output_dir
         logging.info(req)
 
-        return self.detect_one(req)
+        resp = self.detect_one(req)
+        if resp.HasField('localization'):
+            for mir in resp.localization.video_mask:  # type: analytic_pb2.MaskIntRange
+                mir.mask.uri = self.o_unmap(mir.mask.uri)
+            for mir in resp.localization.video_mask_optout:  # type: analytic_pb2.MaskIntRange
+                mir.mask.uri = self.o_unmap(mir.mask.uri)
+
+        for msg in resp.supplement:
+            if msg.uri:
+                msg.uri = self.o_unmap(msg.uri)
+        return resp
 
     def multi_img_manip(self, img, output_dir, resource_kv):
         # type: (str, str, dict) -> analytic_pb2.ImageManipulation
@@ -317,16 +335,25 @@ def multi_img_manip(self, img, output_dir, resource_kv):
         output_dir = self.o_map(output_dir)
         resources = []
         for key, uri in resource_kv.items():
-            resources.append({"key": key, "uri": uri, "type": get_media_type(uri)[0]})
+            resources.append({"key": key, "uri": self.map(uri), "type": get_media_type(uri)[0]})
         req = analytic_pb2.ImageManipulationRequest(options=self.options, resources=resources)
         mime, _ = get_media_type(img)
         req.image.uri = img
         req.image.type = mime
         req.request_id = str(uuid.uuid4())
         req.out_dir = output_dir
-        logging.info(req)
-
-        return self.detect_one(req)
+
+        resp = self.detect_one(req)
+        if resp.localization.mask.uri:
+            resp.localization.mask.uri = self.o_unmap(resp.localization.mask.uri)
+        if resp.localization.mask_optout.uri:
+            resp.localization.mask_optout.uri = self.o_unmap(resp.localization.mask_optout.uri)
+
+        for msg in resp.supplement:
+            if msg.uri:
+                msg.uri = self.o_unmap(msg.uri)
+
+        return resp
 
 
     def detect_batch(self, dir, out, make_dirs=False):