Spaces:

namchain
/

STAR

Sleeping

STAR / app.py

Updated

7668eb6 verified over 1 year ago

4.12 kB

import streamlit as st
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
import json
import cv2
import requests
from io import BytesIO
from src.model import SwinTransformerMultiLabel  # Import from src folder

# Title and description
st.title("STAR Multi-Label Classifier")
st.write("Upload an image or provide a URL to classify and blur sensitive areas.")

# Load class labels from JSON
label_file = "data/labels.json"
with open(label_file, "r") as f:
    label_data = json.load(f)

# Extract unique class labels
class_labels = sorted(set(tag for tags in label_data.values() for tag in tags))
NUM_CLASSES = len(class_labels)

# Define sensitive body parts to blur
sensitive_labels = {"breast", "vagina", "penis"}  # Adjust based on your dataset

# Load trained model
model_path = "models/star.pth"
model = SwinTransformerMultiLabel(num_classes=NUM_CLASSES)
model.load_state_dict(torch.load(model_path, map_location="cpu"))
model.eval()

# Define image preprocessing transformations
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# Function to detect sensitive areas using OpenCV
def detect_sensitive_areas(image):
    """Detect sensitive areas (breasts, vagina, penis) using OpenCV."""
    image_cv = np.array(image)  # Convert PIL to NumPy
    gray = cv2.cvtColor(image_cv, cv2.COLOR_RGB2GRAY)

    # Load OpenCV pre-trained Haarcascade models (ensure these files exist)
    body_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_fullbody.xml")

    # Detect full bodies (basic approach, not perfect for specific areas)
    bodies = body_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

    # Generate a blank mask for blurring
    mask = np.zeros_like(gray)

    # Mark detected regions in mask
    for (x, y, w, h) in bodies:
        mask[y:y+h, x:x+w] = 1  # Mark full-body regions (adjust for more precision)

    return mask

# Function to blur only detected sensitive areas
def blur_sensitive_parts(image, mask, blur_intensity=25):
    """Blurs only the detected sensitive areas while keeping the rest of the image clear."""
    image_np = np.array(image)  # Convert PIL to NumPy

    blurred = cv2.GaussianBlur(image_np, (51, 51), blur_intensity)  # Apply blur
    mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)  # Convert mask to 3 channels
    result = np.where(mask == 1, blurred, image_np)  # Blend only blurred areas

    return Image.fromarray(result)  # Convert back to PIL Image

# UI for image input: Upload or URL
option = st.radio("Choose an input method:", ("Upload Image", "Enter Image URL"))

image = None  # Placeholder for the image

if option == "Upload Image":
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        image = Image.open(uploaded_file).convert("RGB")

elif option == "Enter Image URL":
    image_url = st.text_input("Enter image URL:")
    if image_url:
        try:
            response = requests.get(image_url)
            image = Image.open(BytesIO(response.content)).convert("RGB")
        except Exception as e:
            st.error(f"Error fetching image: {e}")

# Process the image if provided
if image:
    # Preprocess image for model
    img_tensor = transform(image).unsqueeze(0)

    # Perform inference
    with torch.no_grad():
        output = model(img_tensor)
        predicted_indices = [i for i in range(NUM_CLASSES) if output[0][i] > 0.5]
        predicted_labels = [class_labels[i] for i in predicted_indices]

    # Detect sensitive areas
    mask = detect_sensitive_areas(image)

    # Apply selective blurring
    blurred_image = blur_sensitive_parts(image, mask)

    # Display only the blurred image
    st.image(blurred_image, caption="Blurred Image (Sensitive Areas Only)", use_container_width=True)

    # Display results
    st.write("✅ **Predicted Labels:**", ", ".join(predicted_labels) if predicted_labels else "No sensitive content detected")

	import streamlit as st
	import torch
	import numpy as np
	from PIL import Image
	from torchvision import transforms
	import json
	import cv2
	import requests
	from io import BytesIO
	from src.model import SwinTransformerMultiLabel # Import from src folder

	# Title and description
	st.title("STAR Multi-Label Classifier")
	st.write("Upload an image or provide a URL to classify and blur sensitive areas.")

	# Load class labels from JSON
	label_file = "data/labels.json"
	with open(label_file, "r") as f:
	label_data = json.load(f)

	# Extract unique class labels
	class_labels = sorted(set(tag for tags in label_data.values() for tag in tags))
	NUM_CLASSES = len(class_labels)

	# Define sensitive body parts to blur
	sensitive_labels = {"breast", "vagina", "penis"} # Adjust based on your dataset

	# Load trained model
	model_path = "models/star.pth"
	model = SwinTransformerMultiLabel(num_classes=NUM_CLASSES)
	model.load_state_dict(torch.load(model_path, map_location="cpu"))
	model.eval()

	# Define image preprocessing transformations
	transform = transforms.Compose([
	transforms.Resize((224, 224)),
	transforms.ToTensor(),
	transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
	])

	# Function to detect sensitive areas using OpenCV
	def detect_sensitive_areas(image):
	"""Detect sensitive areas (breasts, vagina, penis) using OpenCV."""
	image_cv = np.array(image) # Convert PIL to NumPy
	gray = cv2.cvtColor(image_cv, cv2.COLOR_RGB2GRAY)

	# Load OpenCV pre-trained Haarcascade models (ensure these files exist)
	body_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_fullbody.xml")

	# Detect full bodies (basic approach, not perfect for specific areas)
	bodies = body_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

	# Generate a blank mask for blurring
	mask = np.zeros_like(gray)

	# Mark detected regions in mask
	for (x, y, w, h) in bodies:
	mask[y:y+h, x:x+w] = 1 # Mark full-body regions (adjust for more precision)

	return mask

	# Function to blur only detected sensitive areas
	def blur_sensitive_parts(image, mask, blur_intensity=25):
	"""Blurs only the detected sensitive areas while keeping the rest of the image clear."""
	image_np = np.array(image) # Convert PIL to NumPy

	blurred = cv2.GaussianBlur(image_np, (51, 51), blur_intensity) # Apply blur
	mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2) # Convert mask to 3 channels
	result = np.where(mask == 1, blurred, image_np) # Blend only blurred areas

	return Image.fromarray(result) # Convert back to PIL Image

	# UI for image input: Upload or URL
	option = st.radio("Choose an input method:", ("Upload Image", "Enter Image URL"))

	image = None # Placeholder for the image

	if option == "Upload Image":
	uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
	if uploaded_file is not None:
	image = Image.open(uploaded_file).convert("RGB")

	elif option == "Enter Image URL":
	image_url = st.text_input("Enter image URL:")
	if image_url:
	try:
	response = requests.get(image_url)
	image = Image.open(BytesIO(response.content)).convert("RGB")
	except Exception as e:
	st.error(f"Error fetching image: {e}")

	# Process the image if provided
	if image:
	# Preprocess image for model
	img_tensor = transform(image).unsqueeze(0)

	# Perform inference
	with torch.no_grad():
	output = model(img_tensor)
	predicted_indices = [i for i in range(NUM_CLASSES) if output[0][i] > 0.5]
	predicted_labels = [class_labels[i] for i in predicted_indices]

	# Detect sensitive areas
	mask = detect_sensitive_areas(image)

	# Apply selective blurring
	blurred_image = blur_sensitive_parts(image, mask)

	# Display only the blurred image
	st.image(blurred_image, caption="Blurred Image (Sensitive Areas Only)", use_container_width=True)

	# Display results
	st.write("✅ Predicted Labels:", ", ".join(predicted_labels) if predicted_labels else "No sensitive content detected")