Update app.py

app.py

@@ -206,35 +206,41 @@ def resize_image(image):
     return image
 
 
-def process(image,left_mm, top_mm, right_mm, bottom_mm, ruler, resize,actual_height_in_inches):
+def process(image,left_mm, top_mm, right_mm, bottom_mm, ruler, resize,actual_height_in_inches,background):
 
     # prepare input
     orig_image = Image.fromarray(image)
     print("orig size:",orig_image.size)
     w,h = orig_im_size = orig_image.size
     image = resize_image(orig_image)
-    im_np = np.array(image)
-    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
-    im_tensor = torch.unsqueeze(im_tensor,0)
-    im_tensor = torch.divide(im_tensor,255.0)
-    im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
-    if torch.cuda.is_available():
-        im_tensor=im_tensor.cuda()
-
-    #inference
-    result=net(im_tensor)
-    # post process
-    result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
-    ma = torch.max(result)
-    mi = torch.min(result)
-    result = (result-mi)/(ma-mi)    
-    # image to pil
-    im_array = (result*255).cpu().data.numpy().astype(np.uint8)
-    pil_im = Image.fromarray(np.squeeze(im_array))
-    # paste the mask on the original image
-    new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
-    new_im.paste(orig_image, mask=pil_im)
-    # new_orig_image = orig_image.convert('RGBA')
+    new_im=image;
+
+    if background:
+        # remove background
+        im_np = np.array(image)
+        im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
+        im_tensor = torch.unsqueeze(im_tensor,0)
+        im_tensor = torch.divide(im_tensor,255.0)
+        im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
+        if torch.cuda.is_available():
+            im_tensor=im_tensor.cuda()
+
+        #inference
+        result=net(im_tensor)
+        # post process
+        result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
+        ma = torch.max(result)
+        mi = torch.min(result)
+        result = (result-mi)/(ma-mi)    
+        # image to pil
+        im_array = (result*255).cpu().data.numpy().astype(np.uint8)
+        pil_im = Image.fromarray(np.squeeze(im_array))
+        # paste the mask on the original image
+        new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
+        new_im.paste(orig_image, mask=pil_im)
+        # new_orig_image = orig_image.convert('RGBA')
+    else:
+        new_im=image;
 
     new_im= crop_image(new_im, left_mm, top_mm, right_mm, bottom_mm)
     if ruler:
@@ -298,7 +304,9 @@ demo = gr.Interface(fn=process,
                         gr.Number(label="Bottom Crop (mm)",value=0),
                         gr.Checkbox(label="Ruler!"),
                         gr.Number(label="Resize (between 0.1 and 0.9)",value=0),
-                        gr.Number(label="Plant height in inches",value=0)
+                        gr.Number(label="Plant height in inches",value=0),
+                        gr.Checkbox(label="Remove Background?")
+
                     ],
                     outputs="image", examples=examples, title=title, description=description)