Spaces:
Sleeping
Sleeping
from transformers import BlipProcessor, BlipForConditionalGeneration | |
from sentence_transformers import SentenceTransformer | |
import gradio as gr | |
from PIL import Image | |
import torch | |
import requests | |
from bs4 import BeautifulSoup | |
# Load BLIP Model | |
blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") | |
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base") | |
# Map common fish names to FishBase scientific names | |
name_map = { | |
"pufferfish": "Tetraodon", | |
"stonefish": "Synanceia", | |
"lionfish": "Pterois", | |
"tuna": "Thunnus", | |
"salmon": "Salmo-salar", | |
"catfish": "Ictalurus", | |
"tilapia": "Oreochromis" | |
} | |
# Poisonous species (scientific names) | |
poisonous_species = ["Tetraodon", "Synanceia", "Pterois"] | |
# FishBase scraping function | |
def get_fishbase_summary(scientific_name): | |
search_url = f"https://www.fishbase.se/summary/{scientific_name}.html" | |
try: | |
response = requests.get(search_url, timeout=10) | |
if response.status_code != 200: | |
return f"FishBase entry not found for: {scientific_name}" | |
soup = BeautifulSoup(response.text, "html.parser") | |
summary_section = soup.find("div", {"id": "ssummary"}) | |
if summary_section: | |
paragraphs = summary_section.find_all("p") | |
text = "\n\n".join(p.get_text(strip=True) for p in paragraphs if p.get_text(strip=True)) | |
return text or f"No summary available for {scientific_name}" | |
else: | |
return f"No detailed summary found for {scientific_name}" | |
except Exception as e: | |
return f"Error fetching FishBase data for {scientific_name}: {str(e)}" | |
# Fish identification function | |
def identify_fish(image): | |
# Step 1: Generate caption from image | |
inputs = blip_processor(image, return_tensors="pt") | |
out = blip_model.generate(**inputs) | |
caption = blip_processor.decode(out[0], skip_special_tokens=True) | |
# Step 2: Extract fish name from caption | |
fish_name = None | |
for name in name_map: | |
if name in caption.lower(): | |
fish_name = name | |
break | |
if not fish_name: | |
return f"❌ Could not identify a known fish species in the image caption: '{caption}'" | |
# Step 3: Lookup in FishBase | |
scientific_name = name_map[fish_name] | |
summary = get_fishbase_summary(scientific_name) | |
# Step 4: Check toxicity | |
is_poisonous = "Yes 🧪" if scientific_name in poisonous_species else "No ✅" | |
# Step 5: Final Output | |
return f"**Image Caption:** {caption}\n\n**Detected Fish:** {fish_name.title()}\n**Scientific Name:** {scientific_name}\n**Poisonous:** {is_poisonous}\n\n**📚 FishBase Info:**\n{summary}" | |
# Gradio UI | |
demo = gr.Interface( | |
fn=identify_fish, | |
inputs=gr.Image(type="pil"), | |
outputs="markdown", | |
title="🐟 Smart Fish Identifier (BLIP + FishBase)", | |
description="Upload a fish image. We use BLIP to describe the fish, match it with known species, then fetch info from FishBase to check if it's poisonous." | |
) | |
if __name__ == '__main__': | |
demo.launch() | |