Spaces:
Runtime error
Runtime error
File size: 3,514 Bytes
95bc3ae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 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() |