Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import CLIPModel, CLIPProcessor | |
| from PIL import Image | |
| # Step 1: Load Fine-Tuned Model from Hugging Face Model Hub | |
| model_name = "quadranttechnologies/retail-content-safety-clip-finetuned" | |
| print("Initializing the application...") | |
| try: | |
| print("Loading the model from Hugging Face Model Hub...") | |
| model = CLIPModel.from_pretrained(model_name, trust_remote_code=True) | |
| processor = CLIPProcessor.from_pretrained(model_name) | |
| print("Model and processor loaded successfully.") | |
| except Exception as e: | |
| print(f"Error loading the model or processor: {e}") | |
| raise RuntimeError(f"Failed to load model: {e}") | |
| # Step 2: Define the Inference Function | |
| def classify_image(image): | |
| """ | |
| Classify an image as 'safe' or 'unsafe' and return probabilities. | |
| Args: | |
| image (PIL.Image.Image): Uploaded image. | |
| Returns: | |
| dict: Classification results or an error message. | |
| """ | |
| try: | |
| print("Starting image classification...") | |
| # Validate input | |
| if image is None: | |
| raise ValueError("No image provided. Please upload a valid image.") | |
| # Validate image format | |
| if not hasattr(image, "convert"): | |
| raise ValueError("Invalid image format. Please upload a valid image (JPEG, PNG, etc.).") | |
| # Define categories | |
| categories = ["safe", "unsafe"] | |
| # Process the image with the processor | |
| print("Processing the image...") | |
| inputs = processor(text=categories, images=image, return_tensors="pt", padding=True) | |
| print(f"Processed inputs: {inputs}") | |
| # Run inference with the model | |
| print("Running model inference...") | |
| outputs = model(**inputs) | |
| logits_per_image = outputs.logits_per_image # Image-text similarity scores | |
| print(f"Logits per image: {logits_per_image}") | |
| # Apply softmax to convert logits to probabilities | |
| probs = logits_per_image.softmax(dim=1) # Convert logits to probabilities | |
| print(f"Softmax probabilities: {probs}") | |
| # Extract probabilities for each category | |
| safe_prob = probs[0][0].item() | |
| unsafe_prob = probs[0][1].item() | |
| # Normalize probabilities to ensure they sum to 100% | |
| total_prob = safe_prob + unsafe_prob | |
| safe_percentage = (safe_prob / total_prob) * 100 | |
| unsafe_percentage = (unsafe_prob / total_prob) * 100 | |
| print(f"Normalized percentages: safe={safe_percentage}, unsafe={unsafe_percentage}") | |
| # Return results | |
| return { | |
| "safe": safe_percentage, | |
| "unsafe": unsafe_percentage | |
| } | |
| except Exception as e: | |
| print(f"Error during classification: {e}") | |
| return {"Error": str(e)} | |
| # Step 3: Set Up Gradio Interface | |
| iface = gr.Interface( | |
| fn=classify_image, | |
| inputs=gr.Image(type="pil"), | |
| outputs=gr.Label(num_top_classes=2), # Use gr.Label to display probabilities with a bar-style visualization | |
| title="Content Safety Classification", | |
| description="Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities.", | |
| ) | |
| # Step 4: Launch Gradio Interface | |
| if __name__ == "__main__": | |
| print("Launching the Gradio interface...") | |
| iface.launch() | |