Create app.py

app.py

@@ -0,0 +1,140 @@
+import streamlit as st
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
+# from IPython.display import display, Image
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+from rembg import remove
+from PIL import Image
+
+
+# Content of detectObjects.py file 
+# import detectObjects
+import ultralytics
+from ultralytics import YOLO
+
+model = YOLO('yolov8n.pt')
+sam_checkpoint = "sam_vit_b_01ec64.pth"
+model_type = "vit_b"
+sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+predictor = SamPredictor(sam)
+
+def detected_objects(filename:str):
+    results = model.predict(source=filename, conf=0.25)
+
+    categories = results[0].names
+
+    dc = []
+    for i in range(len(results[0])):
+        cat = results[0].boxes[i].cls
+        dc.append(categories[int(cat)])
+
+    print(dc)
+    return results, dc
+
+def show_mask(mask, ax, random_color=False):
+    if random_color:
+        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
+    else:
+        color = np.array([30/255, 144/255, 255/255, 0.6])
+    h, w = mask.shape[-2:]
+    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+    ax.imshow(mask_image)
+
+def show_points(coords, labels, ax, marker_size=375):
+    pos_points = coords[labels==1]
+    neg_points = coords[labels==0]
+    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
+    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
+
+def show_box(box, ax):
+    x0, y0 = box[0], box[1]
+    w, h = box[2] - box[0], box[3] - box[1]
+    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))
+
+st.title('Extract Objects From Image')
+
+uploaded_file = st.file_uploader('Upload an image')
+
+if uploaded_file is not None:
+    # To read file as bytes:
+    bytes_data = uploaded_file.getvalue()
+    with open('uploaded_file.png','wb') as file:
+        file.write(uploaded_file.getvalue())
+
+    # Detect objects in the uploaded image
+    # results, dc = detectObjects.detected_objects('uploaded_file.png')
+    results, dc = detected_objects('uploaded_file.png')
+
+    st.write(dc)
+
+    option = st.selectbox("Which object would you like to extract?", tuple(dc))
+    # print(option)
+    index_of_the_choosen_detected_object = tuple(dc).index(option)
+
+    if st.button('Extract'):
+        for result in results:
+            boxes = result.boxes
+
+        bbox=boxes.xyxy.tolist()[index_of_the_choosen_detected_object]
+        # sam_checkpoint = "sam_vit_b_01ec64.pth"
+        # model_type = "vit_b"
+        # sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+        # predictor = SamPredictor(sam)
+
+        image = cv2.cvtColor(cv2.imread('uploaded_file.png'), cv2.COLOR_BGR2RGB)
+        predictor.set_image(image)
+
+        input_box = np.array(bbox)
+
+        masks, _, _ = predictor.predict(
+            point_coords=None,
+            point_labels=None,
+            box=input_box[None, :],
+            multimask_output=False,
+        )
+
+        # plt.figure(figsize=(10, 10))
+        # st.image(image)
+        # plt.imshow(image)
+        # show_mask(masks[0], plt.gca())
+        # show_box(input_box, plt.gca())
+        # plt.axis('off')
+        # plt.show()
+
+        segmentation_mask = masks[0]
+        binary_mask = np.where(segmentation_mask > 0.5, 1, 0)
+
+        white_background = np.ones_like(image) * 255
+
+        new_image = white_background * (1 - binary_mask[..., np.newaxis]) + image * binary_mask[..., np.newaxis]
+
+
+        plt.imsave('extracted_image.jpg', new_image.astype(np.uint8))
+        # st.image('extracted_image.jpg')
+
+        # Store path of the image in the variable input_path
+        input_path = 'extracted_image.jpg'
+
+        # Store path of the output image in the variable output_path
+        output_path = 'finalExtracted.png'
+
+        # Processing the image
+        input = Image.open(input_path)
+
+        # Removing the background from the given Image
+        output = remove(input)
+
+        #Saving the image in the given path
+        output.save(output_path)
+        # st.image(output_path)
+
+        with open("finalExtracted.png", "rb") as file:
+            btn = st.download_button(
+                label="Download final image",
+                data=file,
+                file_name="finalExtracted.png",
+                mime="image/png",
+            )
+
+        # bbox=boxes.xyxy.tolist()[0]