Update app.py

app.py

@@ -74,6 +74,10 @@ class Prober:
 
         self.model = self.model.eval()
 
+    def preview_image(self, idx):
+        img_path, target, = self.df.loc[idx][['file_path','bbox']].values
+        img = Image.open(self.zipfile.open(img_path)).convert('RGB')
+        return img
 
     @torch.no_grad()
     def probe(self, idx, re, search_by_sample_id: bool= True):
@@ -82,38 +86,40 @@ class Prober:
         else: 
             img_path, target = self.df[self.df.image_id == idx][['file_path','bbox']].values[0]
         img = Image.open(self.zipfile.open(img_path)).convert('RGB')
-        W0, H0 = img.size
-        sample = {
-            'image': img,
-            'image_size': (H0, W0),  # image original size
-            'bbox': torch.tensor([copy(target)]),
-            'bbox_raw': torch.tensor([copy(target)]),
-            'mask': torch.ones((1, H0, W0), dtype=torch.float32),  # visibiity mask
-            'mask_bbox': None,  # target bbox mask
-        } 
-        sample = self.transform(sample)
-        tok = self.tokenizer(re,
-                             max_length=30,
-                             return_tensors='pt',
-                             truncation=True)
-        inn = {'image': torch.stack([sample['image']]),
-               'mask': torch.stack([sample['mask']]),
-               'tok': tok}
-        output = undo_box_transforms_batch(self.model(inn)[0],
-                                           [sample['tr_param']]).numpy().tolist()[0]
-        img1 = ImageDraw.Draw(img)
-        #img1.rectangle(target, outline ="#0000FF00", width=3)
-        img1.rectangle(output, outline ="#00FF0000", width=3)
-        return img
-
-
+        if re != "":
+            W0, H0 = img.size
+            sample = {
+                'image': img,
+                'image_size': (H0, W0),  # image original size
+                'bbox': torch.tensor([copy(target)]),
+                'bbox_raw': torch.tensor([copy(target)]),
+                'mask': torch.ones((1, H0, W0), dtype=torch.float32),  # visibiity mask
+                'mask_bbox': None,  # target bbox mask
+            } 
+            sample = self.transform(sample)
+            tok = self.tokenizer(re,
+                                 max_length=30,
+                                 return_tensors='pt',
+                                 truncation=True)
+            inn = {'image': torch.stack([sample['image']]),
+                   'mask': torch.stack([sample['mask']]),
+                   'tok': tok}
+            output = undo_box_transforms_batch(self.model(inn)[0],
+                                               [sample['tr_param']]).numpy().tolist()[0]
+            img1 = ImageDraw.Draw(img)
+            #img1.rectangle(target, outline ="#0000FF00", width=3)
+            img1.rectangle(output, outline ="#00FF0000", width=3)
+            return img
+        else:
+            return img
+    
 prober = Prober(
     df_path = 'data/val-sim_metric.json',
     dataset_path = "data/saiapr_tc-12.zip",
     model_checkpoint= "cache/20211220_191132_refclef_32_512_resnet50_8_6_8_0.1_0.1_0.1_0_0.0001_0.0_12_4_90_1_0_0_0/best.ckpt"
 )
 
-demo = gr.Interface(fn=prober.probe, inputs=["number", "text", "checkbox"], outputs="image")
+demo = gr.Interface(fn=prober.probe, inputs=["number", "text"], outputs="image", live=True)
 
 demo.queue(concurrency_count=10)
 demo.launch(debug=True)