Added weights and code fixed

app.py

@@ -1,7 +1,63 @@
 import gradio as gr
+import torch
+import numpy as np
+from transformers import AutoModel
+from transformers import SamModel, SamConfig, SamProcessor
 
-def greet(name):
-    return "Hello " + name + "!"
 
-iface = gr.Interface(fn= greet, inputs="text", outputs="text", title="Greeter")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model_config = SamConfig.from_pretrained("./checkpoint",local_files_only=True)
+processor = SamProcessor.from_pretrained("./checkpoint",local_files_only=True)
+model = SamModel.from_pretrained("./checkpoint",local_files_only=True)
+
+
+
+def get_bbox(gt_map):
+    
+    if np.sum(gt_map) == 0:
+        return [0, 0, gt_map.shape[1], gt_map.shape[0]]
+
+    y_indices, x_indices = np.where(gt_map > 0)
+    x_min, x_max = np.min(x_indices), np.max(x_indices)
+    y_min, y_max = np.min(y_indices), np.max(y_indices)
+
+    H, W = gt_map.shape
+    x_min = max(0, x_min - np.random.randint(0, 20))
+    x_max = min(W, x_max + np.random.randint(0, 20))
+    y_min = max(0, y_min - np.random.randint(0, 20))
+    y_max = min(H, y_max + np.random.randint(0, 20))
+
+    bbox = [x_min,y_min,x_max,y_max]
+
+    return bbox
+
+
+
+def greet(image):
+    
+    image = Image.open(image)
+    image = image.resize((256, 256)) 
+    
+    gt_mask = np.array(image)
+    prompt = get_bbox(gt_mask)
+
+    inputs = processor(images=image, input_boxes=[[prompt]], return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    model.eval()
+    with torch.no_grad():
+        outputs = model(**inputs,multimask_outputs=False)
+    
+    # outputs = outputs.logits[0].cpu().numpy()
+    # outputs = np.argmax(outputs, axis=0)
+    # outputs = Image.fromarray(outputs)
+    # return outputs
+    seg_prob = torch.sigmoid(outputs.pred_masks.squeeze(0))
+    seg_prob = seg_prob.cpu().numpy().squeeze()
+    seg_prob = (seg_prob > 0.5).astype(np.uint8)
+
+    return seg_prob
+
+
+iface = gr.Interface(fn= greet, inputs="image", outputs="image", title="Greeter")
 iface.launch()

checkpoint/config.json

@@ -0,0 +1,23 @@
+{
+  "_name_or_path": "./MainDir/Untitled Folder/checkpoint_sam_torch",
+  "architectures": [
+    "SamModel"
+  ],
+  "initializer_range": 0.02,
+  "mask_decoder_config": {
+    "model_type": ""
+  },
+  "model_type": "sam",
+  "prompt_encoder_config": {
+    "model_type": ""
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.37.0",
+  "vision_config": {
+    "dropout": 0.0,
+    "initializer_factor": 1.0,
+    "intermediate_size": 6144,
+    "model_type": "",
+    "projection_dim": 512
+  }
+}

checkpoint/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a481b414a5dbe9cba508256e6100dba23eb4e9225f132add55176d51e102215
+size 374979376

checkpoint/preprocessor_config.json

@@ -0,0 +1,28 @@
+{
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_pad": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_processor_type": "SamImageProcessor",
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "pad_size": {
+    "height": 1024,
+    "width": 1024
+  },
+  "processor_class": "SamProcessor",
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "longest_edge": 1024
+  }
+}

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+torch
+numpy
+transformers
+Pillow