v1

app.py

@@ -0,0 +1,110 @@
+import torch
+import numpy as np
+import cv2
+from transformers import SamModel, SamProcessor
+import gradio as gr
+
+# Load the SAM model and processor from Hugging Face
+model_id = "facebook/sam-vit-huge"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+sam = SamModel.from_pretrained(model_id).to(device)
+processor = SamProcessor.from_pretrained(model_id)
+
+def segment_rocks(image):
+    # Preprocess the image
+    inputs = processor(image, return_tensors="pt").to(device)
+    
+    # Generate image embeddings
+    with torch.no_grad():
+        image_embeddings = sam.get_image_embeddings(inputs["pixel_values"])
+    
+    # Generate masks
+    masks = []
+    for i in range(3):  # Generate multiple masks
+        inputs = processor(
+            image,
+            input_points=None,
+            return_tensors="pt",
+            input_boxes=[[[0, 0, image.shape[1], image.shape[0]]]],
+        ).to(device)
+        
+        with torch.no_grad():
+            outputs = sam(
+                input_points=inputs["input_points"],
+                input_boxes=inputs["input_boxes"],
+                image_embeddings=image_embeddings,
+                multimask_output=True,
+            )
+        
+        masks.extend(outputs.pred_masks.squeeze().cpu().numpy())
+    
+    return masks
+
+def compute_rock_properties(mask):
+    # Find contours of the mask
+    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    
+    properties = []
+    for contour in contours:
+        # Compute area
+        area = cv2.contourArea(contour)
+        
+        # Compute perimeter
+        perimeter = cv2.arcLength(contour, True)
+        
+        # Compute circularity
+        circularity = 4 * np.pi * area / (perimeter ** 2) if perimeter > 0 else 0
+        
+        # Fit an ellipse to get major and minor axes
+        if len(contour) >= 5:
+            ellipse = cv2.fitEllipse(contour)
+            major_axis = max(ellipse[1])
+            minor_axis = min(ellipse[1])
+            aspect_ratio = major_axis / minor_axis if minor_axis > 0 else 0
+        else:
+            major_axis = minor_axis = aspect_ratio = 0
+        
+        properties.append({
+            'area': area,
+            'perimeter': perimeter,
+            'circularity': circularity,
+            'major_axis': major_axis,
+            'minor_axis': minor_axis,
+            'aspect_ratio': aspect_ratio
+        })
+    
+    return properties
+
+def process_image(input_image):
+    # Convert to RGB if needed
+    if input_image.shape[2] == 4:  # RGBA
+        input_image = cv2.cvtColor(input_image, cv2.COLOR_RGBA2RGB)
+    elif len(input_image.shape) == 2:  # Grayscale
+        input_image = cv2.cvtColor(input_image, cv2.COLOR_GRAY2RGB)
+    
+    masks = segment_rocks(input_image)
+    
+    results = []
+    for i, mask in enumerate(masks):
+        properties = compute_rock_properties(mask)
+        
+        # Visualize the segmentation
+        masked_image = input_image.copy()
+        masked_image[mask] = (masked_image[mask] * 0.7 + np.array([255, 0, 0]) * 0.3).astype(np.uint8)
+        
+        results.append((masked_image, f"Rock {i+1} properties: {properties}"))
+    
+    return results
+
+# Gradio interface
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="numpy"),
+    outputs=[gr.Image(type="numpy"), gr.Textbox(label="Properties")] * 3,
+    title="Rock Segmentation using SAM",
+    description="Upload an image to segment rocks and compute their properties."
+)
+
+# Launch the interface
+iface.launch()

requirements.txt

@@ -0,0 +1,6 @@
+transformers==4.30.0
+torch==1.9.0
+torchvision==0.10.0
+opencv-python==4.5.3.56
+numpy==1.21.0
+gradio==3.35.2