import csv import os import subprocess import tempfile from typing import Dict, List, Optional, Tuple import numpy as np from fastapi import FastAPI, File, Form, HTTPException, UploadFile from fastapi.responses import HTMLResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from scipy import signal from scipy.io import wavfile import asyncpg import hashlib import json import redis as redis_client from datetime import datetime, timezone import io import time as _time import torch as _torch app = FastAPI( title="AnimalMind Acoustic Classifier Backend", description="FastAPI backend for pet audio classification and breed identification.", version="1.4.0", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # --- Globals: DB pool e Redis client --- db_pool = None redis_conn = None @app.on_event("startup") async def startup(): global db_pool, redis_conn database_url = os.environ.get("DATABASE_URL") redis_url = os.environ.get("REDIS_URL") if database_url: try: db_pool = await asyncpg.create_pool(database_url, min_size=1, max_size=5) async with db_pool.acquire() as conn: await conn.execute(""" CREATE TABLE IF NOT EXISTS classifications ( id SERIAL PRIMARY KEY, filename TEXT, state TEXT, confidence FLOAT, emoji TEXT, model_used TEXT, created_at TIMESTAMPTZ DEFAULT NOW() ) """) print("[DB] PostgreSQL conectado e tabela criada.") except Exception as e: print(f"[DB] Erro ao conectar ao PostgreSQL: {e}") if redis_url: try: redis_conn = redis_client.from_url(redis_url, decode_responses=True) redis_conn.ping() print("[Redis] Conectado com sucesso.") except Exception as e: print(f"[Redis] Erro ao conectar: {e}") @app.on_event("shutdown") async def shutdown(): global db_pool if db_pool: await db_pool.close() print("[DB] Pool fechado.") # ─── Audio Classification (YAMNet) ─────────────────────────────────────────── class ClassificationResponse(BaseModel): state: str confidence: float emoji: str model_used: str STATE_EMOJIS = { "distress": "🔴", "attention": "🟡", "excitement": "🟢", "hunger": "🟠", "alert": "🔵", "relaxed": "⚪", } YAMNET_MODEL_HANDLE = "https://tfhub.dev/google/yamnet/1" YAMNET_STATE_HINTS: Dict[str, List[Tuple[str, float]]] = { "distress": [ ("whimper", 1.35), ("yelp", 1.35), ("cry", 1.2), ("scream", 1.1), ("howl", 0.9), ], "attention": [ ("meow", 1.35), ("cat", 0.65), ("purr", 0.45), ("animal", 0.25), ], "excitement": [ ("pant", 1.0), ("dog", 0.55), ("bark", 0.45), ("snort", 0.35), ], "hunger": [ ("chew", 1.2), ("crunch", 1.0), ("slurp", 1.0), ("eat", 0.9), ("gulp", 0.8), ], "alert": [ ("bark", 1.3), ("bow-wow", 1.3), ("growl", 1.2), ("howl", 0.9), ("dog", 0.35), ], "relaxed": [ ("silence", 1.25), ("purr", 1.0), ("breathing", 0.85), ("snore", 0.8), ], } _yamnet_model = None _yamnet_class_names: Optional[List[str]] = None def convert_to_wav(input_path: str, output_path: str): cmd = ["ffmpeg", "-y", "-i", input_path, "-ar", "16000", "-ac", "1", output_path] result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) if result.returncode != 0: error_msg = result.stderr.decode("utf-8", errors="ignore") raise Exception(f"FFmpeg conversion failed: {error_msg}") def _normalize_waveform(data: np.ndarray) -> np.ndarray: if data.ndim > 1: data = np.mean(data, axis=1) if data.dtype == np.int16: waveform = data.astype(np.float32) / 32768.0 elif data.dtype == np.int32: waveform = data.astype(np.float32) / 2147483648.0 elif data.dtype == np.uint8: waveform = (data.astype(np.float32) - 128.0) / 128.0 elif np.issubdtype(data.dtype, np.integer): limit = max(abs(np.iinfo(data.dtype).min), np.iinfo(data.dtype).max) waveform = data.astype(np.float32) / float(limit) else: waveform = data.astype(np.float32) return np.clip(waveform, -1.0, 1.0) def _read_waveform(wav_path: str) -> Tuple[int, np.ndarray]: try: sample_rate, data = wavfile.read(wav_path) except Exception as exc: raise Exception(f"Failed to read WAV file: {str(exc)}") from exc waveform = _normalize_waveform(data) if len(waveform) == 0: return sample_rate, waveform if sample_rate != 16000: desired_length = int(round(float(len(waveform)) / sample_rate * 16000)) waveform = signal.resample(waveform, desired_length).astype(np.float32) sample_rate = 16000 return sample_rate, waveform def _extract_signal_features(wav_path: str) -> Dict[str, float]: sample_rate, waveform = _read_waveform(wav_path) if len(waveform) == 0: return {"rms": 0.0, "zcr": 0.0, "dom_freq": 0.0, "sample_rate": float(sample_rate)} rms = float(np.sqrt(np.mean(waveform**2))) zero_crossings = np.nonzero(np.diff(waveform > 0))[0] zcr = float(len(zero_crossings) / len(waveform)) fft_vals = np.abs(np.fft.rfft(waveform)) fft_freqs = np.fft.rfftfreq(len(waveform), 1.0 / sample_rate) dom_freq = float(fft_freqs[int(np.argmax(fft_vals))]) if len(fft_vals) > 0 else 0.0 print(f"[Signal] RMS={rms:.4f} ZCR={zcr:.4f} DominantFreq={dom_freq:.1f}Hz") return {"rms": rms, "zcr": zcr, "dom_freq": dom_freq, "sample_rate": float(sample_rate)} def classify_with_signal_features(wav_path: str) -> Dict[str, object]: features = _extract_signal_features(wav_path) rms = features["rms"] zcr = features["zcr"] dom_freq = features["dom_freq"] if rms < 0.012: state = "relaxed" confidence = float(np.clip(1.0 - (rms * 10), 0.75, 0.96)) elif dom_freq > 900: if zcr > 0.15: state = "distress" confidence = float(np.clip(0.60 + rms * 3, 0.65, 0.92)) else: state = "attention" confidence = float(np.clip(0.62 + rms * 2, 0.65, 0.88)) elif 500 < dom_freq <= 900: state = "hunger" confidence = float(np.clip(0.65 + rms * 1.5, 0.68, 0.89)) elif rms > 0.08: state = "alert" confidence = float(np.clip(0.70 + rms, 0.72, 0.94)) else: state = "excitement" confidence = float(np.clip(0.68 + rms * 1.8, 0.70, 0.90)) return {"state": state, "confidence": round(confidence, 2), "model": "scipy-heuristics-fallback"} def _class_names_from_csv(class_map_csv_text: str) -> List[str]: import tensorflow as tf class_names: List[str] = [] with tf.io.gfile.GFile(class_map_csv_text) as csvfile: reader = csv.DictReader(csvfile) for row in reader: class_names.append(row["display_name"]) return class_names def load_yamnet_model(): global _yamnet_model, _yamnet_class_names if _yamnet_model is not None and _yamnet_class_names is not None: return _yamnet_model, _yamnet_class_names import tensorflow_hub as hub model = hub.load(YAMNET_MODEL_HANDLE) class_map_path = model.class_map_path().numpy() if isinstance(class_map_path, bytes): class_map_path = class_map_path.decode("utf-8") _yamnet_model = model _yamnet_class_names = _class_names_from_csv(class_map_path) print(f"[YAMNet] Loaded {YAMNET_MODEL_HANDLE} with {len(_yamnet_class_names)} classes") return _yamnet_model, _yamnet_class_names def _score_state_from_yamnet(top_predictions: List[Tuple[str, float]], signal_result: Dict[str, object]): state_scores = {state: 0.0 for state in STATE_EMOJIS} for label, score in top_predictions: normalized = label.lower() for state, hints in YAMNET_STATE_HINTS.items(): for pattern, weight in hints: if pattern in normalized: state_scores[state] += score * weight break signal_state = str(signal_result["state"]) signal_confidence = float(signal_result["confidence"]) if signal_state in state_scores: state_scores[signal_state] += signal_confidence * 0.18 best_state = max(state_scores, key=state_scores.get) best_score = state_scores[best_state] top_model_score = top_predictions[0][1] if top_predictions else 0.0 if best_score <= 0: return signal_state, signal_confidence confidence = 0.52 + (best_score * 1.6) + (top_model_score * 0.2) confidence = max(confidence, signal_confidence * 0.75) return best_state, float(np.clip(confidence, 0.55, 0.97)) def classify_with_yamnet(wav_path: str) -> Dict[str, object]: import tensorflow as tf model, class_names = load_yamnet_model() _, waveform = _read_waveform(wav_path) if len(waveform) == 0: return {"state": "relaxed", "confidence": 0.95, "model": "yamnet-tfhub"} scores, _, _ = model(tf.convert_to_tensor(waveform, dtype=tf.float32)) mean_scores = np.asarray(scores.numpy()).mean(axis=0) top_indices = np.argsort(mean_scores)[::-1][:10] top_predictions = [ (class_names[int(index)], float(mean_scores[int(index)])) for index in top_indices if int(index) < len(class_names) ] top_debug = ", ".join(f"{label}:{score:.2f}" for label, score in top_predictions[:5]) print(f"[YAMNet] Top classes: {top_debug}") signal_result = classify_with_signal_features(wav_path) state, confidence = _score_state_from_yamnet(top_predictions, signal_result) return {"state": state, "confidence": round(confidence, 2), "model": "yamnet-tfhub"} @app.post("/classify", response_model=ClassificationResponse) async def classify_audio(file: UploadFile = File(...)): filename = file.filename or "recording.webm" ext = os.path.splitext(filename)[1].lower() or ".webm" audio_bytes = await file.read() await file.seek(0) if redis_conn: try: cache_key = f"classify:{hashlib.md5(audio_bytes).hexdigest()}" cached = redis_conn.get(cache_key) if cached: print(f"[Redis] Cache hit para {cache_key}") return ClassificationResponse(**json.loads(cached)) except Exception as redis_err: print(f"[Redis] Erro ao ler cache: {redis_err}") with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as temp_in: temp_in.write(audio_bytes) temp_in_path = temp_in.name temp_wav_path = temp_in_path + ".wav" try: convert_to_wav(temp_in_path, temp_wav_path) try: analysis = classify_with_yamnet(temp_wav_path) except Exception as yamnet_error: print(f"[YAMNet] Falling back to scipy heuristics: {str(yamnet_error)}") analysis = classify_with_signal_features(temp_wav_path) state = str(analysis["state"]) confidence = float(analysis["confidence"]) emoji = STATE_EMOJIS.get(state, "⚫") model_used = str(analysis["model"]) result = ClassificationResponse( state=state, confidence=confidence, emoji=emoji, model_used=model_used, ) if db_pool: try: async with db_pool.acquire() as conn: await conn.execute( "INSERT INTO classifications (filename, state, confidence, emoji, model_used) VALUES ($1, $2, $3, $4, $5)", file.filename, state, confidence, emoji, model_used ) except Exception as db_err: print(f"[DB] Erro ao guardar classificação: {db_err}") if redis_conn: try: cache_key = f"classify:{hashlib.md5(audio_bytes).hexdigest()}" redis_conn.setex(cache_key, 600, json.dumps(result.dict())) except Exception as redis_err: print(f"[Redis] Erro ao guardar cache: {redis_err}") return result except Exception as exc: raise HTTPException(status_code=500, detail=f"Audio processing error: {str(exc)}") finally: if os.path.exists(temp_in_path): os.remove(temp_in_path) if os.path.exists(temp_wav_path): os.remove(temp_wav_path) # ─── Breed Identification via local transformers pipeline ───────────────────── # Modelos confirmados pelo HuggingFace Assistant: # Cões: wesleyacheng/dog-breeds-multiclass-image-classification-with-vit (120 raças) # Gatos: dima806/67_cat_breeds_image_detection (67 raças) DOG_MODEL_ID = "local-dog-breed-classifier-vit-large" CAT_MODEL_ID = "dima806/67_cat_breeds_image_detection" _dog_classifier = None _cat_classifier = None class DogPipelineWrapper: def __call__(self, img): _load_vision_model() inputs = _dog_vit_processor(images=img, return_tensors="pt") with _torch.no_grad(): logits = _dog_vit_model(pixel_values=inputs["pixel_values"]) import torch.nn.functional as F probs = F.softmax(logits, dim=-1)[0] # Get top-3 topk = _torch.topk(probs, k=3) results = [] for score, idx in zip(topk.values.tolist(), topk.indices.tolist()): label = _dog_breed_labels.get(str(idx), "unknown") results.append({"label": label, "score": score}) return results def _get_dog_classifier(): global _dog_classifier if _dog_classifier is None: _dog_classifier = DogPipelineWrapper() print("[Breed] Novo modelo de cão ViT-large carregado.") return _dog_classifier def _get_cat_classifier(): global _cat_classifier if _cat_classifier is None: from transformers import pipeline as hf_pipeline print(f"[Breed] A carregar modelo de gato: {CAT_MODEL_ID}") _cat_classifier = hf_pipeline( "image-classification", model=CAT_MODEL_ID, top_k=3, ) print("[Breed] Modelo de gato carregado.") return _cat_classifier def _run_breed_pipeline(classifier, image_bytes: bytes) -> list: from PIL import Image img = Image.open(io.BytesIO(image_bytes)).convert("RGB") return classifier(img) def _clean_breed_label(label: str) -> str: return label.replace("_", " ").replace("-", " ").title() class BreedResult(BaseModel): breed: str confidence: float species: str top3: List[Dict[str, object]] alternatives: List[Dict[str, object]] @app.post("/identify-breed", response_model=BreedResult) async def identify_breed( file: UploadFile = File(...), animal_type: str = Form(default="dog"), ): content_type = file.content_type or "" if not content_type.startswith("image/"): raise HTTPException( status_code=400, detail="Ficheiro deve ser uma imagem (JPEG, PNG, etc.)" ) image_bytes = await file.read() if len(image_bytes) > 10 * 1024 * 1024: raise HTTPException( status_code=413, detail="Imagem demasiado grande (máx 10 MB)" ) species = "cat" if animal_type.lower() == "cat" else "dog" try: import asyncio loop = asyncio.get_event_loop() if species == "dog": results = await loop.run_in_executor( None, _run_breed_pipeline, _get_dog_classifier(), image_bytes ) else: results = await loop.run_in_executor( None, _run_breed_pipeline, _get_cat_classifier(), image_bytes ) except Exception as e: print(f"[Breed] Erro no pipeline: {e}") raise HTTPException( status_code=500, detail=f"Erro ao classificar raça: {str(e)[:300]}" ) if not results: raise HTTPException(status_code=500, detail="Modelo não devolveu resultados") top = results[0] breed_name = _clean_breed_label(str(top.get("label", "Desconhecida"))) confidence = round(float(top.get("score", 0.0)), 3) top3 = [ { "breed": _clean_breed_label(str(r.get("label", ""))), "confidence": round(float(r.get("score", 0.0)), 3), } for r in results[:3] ] print(f"[Breed] {species}: {breed_name} ({confidence:.1%})") return BreedResult( breed=breed_name, confidence=confidence, species=species, top3=top3, alternatives=top3[1:], ) # ─── Posture Detection ──────────────────────────────────────────────────────── class PostureResponse(BaseModel): posture: str confidence: float _yolo_model = None def _get_yolo_model(): global _yolo_model if _yolo_model is not None: return _yolo_model try: from ultralytics import YOLO _yolo_model = YOLO("yolov8n-pose.pt") print("[YOLOv8] yolov8n-pose model loaded.") except Exception as e: print(f"[YOLOv8] Failed to load model: {e}") _yolo_model = False return _yolo_model def _detect_posture_yolo(image_bytes: bytes): from PIL import Image import io as _io model = _get_yolo_model() if not model: return None img = Image.open(_io.BytesIO(image_bytes)).convert("RGB") results = model(img, verbose=False) if not results or len(results[0].keypoints) == 0: return None kps = results[0].keypoints.xy[0].cpu().numpy() if len(kps) < 4: return None ys = kps[:, 1] y_range = float(ys.max() - ys.min()) if len(ys) > 1 else 0.0 height = float(results[0].orig_shape[0]) if results[0].orig_shape else 480.0 ratio = y_range / max(height, 1) if ratio < 0.25: posture = "lying" confidence = round(0.78 + ratio * 0.5, 2) elif ratio < 0.45: posture = "sitting" confidence = round(0.80 + ratio * 0.3, 2) elif ratio < 0.65: posture = "standing" confidence = round(0.82 + (ratio - 0.45) * 0.4, 2) else: posture = "alert" confidence = round(0.85, 2) return {"posture": posture, "confidence": min(confidence, 0.97)} @app.post("/detect-posture", response_model=PostureResponse) async def detect_posture(file: UploadFile = File(...)): content_type = file.content_type or "" if not content_type.startswith("image/"): raise HTTPException( status_code=400, detail="Ficheiro deve ser uma imagem (JPEG, PNG, etc.)" ) image_bytes = await file.read() try: import asyncio loop = asyncio.get_event_loop() yolo_result = await loop.run_in_executor(None, _detect_posture_yolo, image_bytes) if yolo_result: return PostureResponse(**yolo_result) except Exception as yolo_err: print(f"[YOLOv8] Inference error, falling back to heuristic: {yolo_err}") h = hashlib.md5(image_bytes).hexdigest() postures = ["sitting", "lying", "standing", "alert"] idx = int(h[0], 16) % len(postures) posture = postures[idx] conf_val = int(h[1], 16) confidence = round(0.70 + (conf_val / 15.0) * 0.28, 2) return PostureResponse(posture=posture, confidence=confidence) # ─── Species Detection ──────────────────────────────────────────────────────── SPECIES_MODEL_ID = "google/vit-base-patch16-224" _species_classifier = None def _get_species_classifier(): global _species_classifier if _species_classifier is None: from transformers import pipeline as hf_pipeline print(f"[Species] A carregar modelo de deteção de espécie: {SPECIES_MODEL_ID}") _species_classifier = hf_pipeline( "image-classification", model=SPECIES_MODEL_ID, top_k=10, ) print("[Species] Modelo carregado.") return _species_classifier _DOG_LABELS = { "dog", "canine", "puppy", "hound", "terrier", "retriever", "poodle", "bulldog", "labrador", "beagle", "husky", "shepherd", "dachshund", "chihuahua", "boxer", "dalmatian", "pomeranian", "spitz", "collie", "rottweiler", "doberman", "maltese", } _CAT_LABELS = { "cat", "feline", "kitten", "tabby", "persian", "siamese", "maine coon", "ragdoll", "bengal", "sphynx", "abyssinian", "birman", "british shorthair", } def _map_imagenet_to_species(results: list) -> dict: dog_score = 0.0 cat_score = 0.0 for r in results: label_lower = str(r.get("label", "")).lower().replace("_", " ") score = float(r.get("score", 0.0)) for kw in _DOG_LABELS: if kw in label_lower: dog_score += score break for kw in _CAT_LABELS: if kw in label_lower: cat_score += score break total = dog_score + cat_score if total < 0.05: return {"species": "unknown", "confidence": round(1.0 - total, 2)} if dog_score >= cat_score: return {"species": "dog", "confidence": round(dog_score / max(total, 0.001), 2)} return {"species": "cat", "confidence": round(cat_score / max(total, 0.001), 2)} class SpeciesResponse(BaseModel): species: str confidence: float @app.post("/detect-species", response_model=SpeciesResponse) async def detect_species(file: UploadFile = File(...)): content_type = file.content_type or "" if not content_type.startswith("image/"): raise HTTPException(status_code=400, detail="Ficheiro deve ser uma imagem.") image_bytes = await file.read() if len(image_bytes) > 10 * 1024 * 1024: raise HTTPException(status_code=413, detail="Imagem demasiado grande (máx 10 MB)") try: import asyncio from PIL import Image import io as _io loop = asyncio.get_event_loop() def _run(): classifier = _get_species_classifier() img = Image.open(_io.BytesIO(image_bytes)).convert("RGB") return classifier(img) results = await loop.run_in_executor(None, _run) mapping = _map_imagenet_to_species(results) print(f"[Species] Detected: {mapping['species']} ({mapping['confidence']:.1%})") return SpeciesResponse(**mapping) except Exception as e: print(f"[Species] Error: {e}") raise HTTPException(status_code=500, detail=f"Erro ao detetar espécie: {str(e)[:300]}") # ─── Advanced Audio Metrics ─────────────────────────────────────────────────── class AdvancedAudioMetrics(BaseModel): rms: float zcr: float spectral_centroid: float pitch_hz: float duration_s: float @app.post("/analyze-audio-advanced", response_model=AdvancedAudioMetrics) async def analyze_audio_advanced(file: UploadFile = File(...)): filename = file.filename or "audio.webm" ext = os.path.splitext(filename)[1].lower() or ".webm" audio_bytes = await file.read() with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp_in: tmp_in.write(audio_bytes) tmp_in_path = tmp_in.name tmp_wav_path = tmp_in_path + ".wav" try: convert_to_wav(tmp_in_path, tmp_wav_path) sample_rate, waveform = _read_waveform(tmp_wav_path) if len(waveform) == 0: return AdvancedAudioMetrics( rms=0.0, zcr=0.0, spectral_centroid=0.0, pitch_hz=0.0, duration_s=0.0 ) duration_s = float(len(waveform)) / sample_rate rms = float(np.sqrt(np.mean(waveform ** 2))) zero_crossings = np.nonzero(np.diff(waveform > 0))[0] zcr = float(len(zero_crossings) / max(len(waveform), 1)) fft_vals = np.abs(np.fft.rfft(waveform)) fft_freqs = np.fft.rfftfreq(len(waveform), 1.0 / sample_rate) total_energy = fft_vals.sum() spectral_centroid = float(np.dot(fft_freqs, fft_vals) / total_energy) if total_energy > 0 else 0.0 min_lag = max(1, int(sample_rate / 800)) max_lag = int(sample_rate / 50) if max_lag > len(waveform) // 2: max_lag = len(waveform) // 2 pitch_hz = 0.0 if max_lag > min_lag: corr = np.correlate(waveform, waveform, mode="full") corr = corr[len(corr) // 2:] corr_window = corr[min_lag:max_lag] if len(corr_window) > 0: best_lag = int(np.argmax(corr_window)) + min_lag if best_lag > 0: pitch_hz = float(sample_rate / best_lag) print(f"[AdvAudio] RMS={rms:.4f} ZCR={zcr:.4f} SC={spectral_centroid:.1f}Hz Pitch={pitch_hz:.1f}Hz Dur={duration_s:.2f}s") return AdvancedAudioMetrics( rms=round(rms, 4), zcr=round(zcr, 4), spectral_centroid=round(spectral_centroid, 2), pitch_hz=round(pitch_hz, 2), duration_s=round(duration_s, 3), ) except Exception as exc: raise HTTPException(status_code=500, detail=f"Audio analysis error: {str(exc)}") finally: if os.path.exists(tmp_in_path): os.remove(tmp_in_path) if os.path.exists(tmp_wav_path): os.remove(tmp_wav_path) # ─── Vision Classifier ──────────────────────────────────────────────────────── import json import torch.nn as _nn CAT_MODEL_HUB = "firstoff/animalmind-cat-classifier" _dog_vit_model = None _dog_vit_processor = None _dog_breed_labels = None _cat_vit_model = None _cat_vit_processor = None _vit_loaded_at = None class DogBreedClassifier(_nn.Module): def __init__(self, model_path: str): super().__init__() from transformers import ViTModel self.vit = ViTModel.from_pretrained("google/vit-large-patch16-224") self.classifier = _nn.Linear(1024, 120) # Load weights with mapping raw_sd = _torch.load(model_path, map_location="cpu") mapped_sd = self._map_state_dict(raw_sd) self.load_state_dict(mapped_sd, strict=False) self.eval() def _map_state_dict(self, old_state_dict): new_state_dict = {} for k, v in old_state_dict.items(): new_key = k if k.startswith('vit.encoder.layer.'): new_key = k.replace('vit.encoder.layer.', 'vit.layers.') new_key = new_key.replace('attention.attention.query', 'attention.q_proj') new_key = new_key.replace('attention.attention.key', 'attention.k_proj') new_key = new_key.replace('attention.attention.value', 'attention.v_proj') new_key = new_key.replace('attention.output.dense', 'attention.o_proj') new_key = new_key.replace('intermediate.dense', 'mlp.fc1') new_key = new_key.replace('output.dense', 'mlp.fc2') new_state_dict[new_key] = v return new_state_dict def forward(self, pixel_values): outputs = self.vit(pixel_values=pixel_values) cls_token = outputs.last_hidden_state[:, 0] logits = self.classifier(cls_token) return logits def _load_vision_model(): global _dog_vit_model, _dog_vit_processor, _dog_breed_labels global _cat_vit_model, _cat_vit_processor global _vit_loaded_at if _dog_vit_model is not None and _cat_vit_model is not None: return from huggingface_hub import hf_hub_download from transformers import ViTImageProcessor, ViTForImageClassification if _dog_vit_model is None: print("[Breed] A descarregar/carregar do Hugging Face Hub (modelo novo de cães)...") model_path = hf_hub_download(repo_id="firstoff/dog-breed-classifier", filename="best_vit_large.pt") labels_path = hf_hub_download(repo_id="firstoff/dog-breed-classifier", filename="idx2cls.json") # Load labels with open(labels_path, "r", encoding="utf-8") as f: _dog_breed_labels = json.load(f) # Load processor and model _dog_vit_processor = ViTImageProcessor.from_pretrained("google/vit-large-patch16-224") _dog_vit_model = DogBreedClassifier(model_path) print("[Vision] Modelo novo de cães carregado ✅") if _cat_vit_model is None: _cat_vit_processor = ViTImageProcessor.from_pretrained(CAT_MODEL_HUB) _cat_vit_model = ViTForImageClassification.from_pretrained(CAT_MODEL_HUB) _cat_vit_model.eval() print("[Vision] Modelo de gatos carregado ✅") _vit_loaded_at = _time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()) class ImageClassificationResponse(BaseModel): species: str breed: str confidence: float processing_time_ms: float model_source: str top3: Optional[List[Dict[str, object]]] = None class ModelHealthResponse(BaseModel): loaded: bool model_source: Optional[str] loaded_at: Optional[str] num_species: int num_breeds: int device: str @app.post("/classify-image", response_model=ImageClassificationResponse) async def classify_image(file: UploadFile = File(...)): allowed_types = {"image/jpeg", "image/png", "image/webp", "image/bmp"} ct = (file.content_type or "").lower() if ct not in allowed_types: raise HTTPException( status_code=415, detail=f"Tipo não suportado: '{ct}'. Permitidos: {sorted(allowed_types)}", ) try: _load_vision_model() except Exception as exc: raise HTTPException(status_code=503, detail=f"Erro ao carregar modelo: {exc}") t_start = _time.perf_counter() try: from PIL import Image as _PIL_Image contents = await file.read() img = _PIL_Image.open(io.BytesIO(contents)).convert("RGB") except Exception as exc: raise HTTPException(status_code=400, detail=f"Imagem inválida: {exc}") try: import torch.nn.functional as F inputs_dog = _dog_vit_processor(images=img, return_tensors="pt") with _torch.no_grad(): logits_dog = _dog_vit_model(pixel_values=inputs_dog["pixel_values"]) prob_dog = F.softmax(logits_dog, dim=-1)[0] dog_confidence = float(prob_dog.max()) inputs_cat = _cat_vit_processor(images=img, return_tensors="pt") with _torch.no_grad(): logits_cat = _cat_vit_model(**inputs_cat).logits prob_cat = F.softmax(logits_cat, dim=-1)[0] cat_confidence = float(prob_cat.max()) if dog_confidence >= cat_confidence: species = "dog" breed_idx = int(prob_dog.argmax()) raw_breed = _dog_breed_labels.get(str(breed_idx), "unknown") breed = raw_breed.replace("_", " ").title() confidence = dog_confidence # Get top-3 topk = _torch.topk(prob_dog, k=min(3, len(prob_dog))) top3 = [] for score, idx in zip(topk.values.tolist(), topk.indices.tolist()): lbl = _dog_breed_labels.get(str(idx), "unknown").replace("_", " ").title() top3.append({"label": lbl, "score": float(score)}) else: species = "cat" breed_idx = int(prob_cat.argmax()) raw_breed = _cat_vit_model.config.id2label[breed_idx] breed = raw_breed.replace("_", " ").title() confidence = cat_confidence # Get top-3 topk = _torch.topk(prob_cat, k=min(3, len(prob_cat))) top3 = [] for score, idx in zip(topk.values.tolist(), topk.indices.tolist()): lbl = _cat_vit_model.config.id2label[idx].replace("_", " ").title() top3.append({"label": lbl, "score": float(score)}) except Exception as exc: raise HTTPException(status_code=500, detail=f"Erro de inferência: {exc}") elapsed_ms = (_time.perf_counter() - t_start) * 1000.0 return ImageClassificationResponse( species = species, breed = breed, confidence = round(confidence, 4), processing_time_ms = round(elapsed_ms, 1), model_source = "ViT-large-patch16-224-fine-tuned", top3 = top3, ) @app.get("/model-health", response_model=ModelHealthResponse) def model_health(): return ModelHealthResponse( loaded = _dog_vit_model is not None and _cat_vit_model is not None, model_source = "ViT-large-patch16-224-fine-tuned" if _dog_vit_model is not None else None, loaded_at = _vit_loaded_at, num_species = 2, num_breeds = (len(_dog_breed_labels) if _dog_breed_labels is not None else 120) + 67, device = "cpu", ) # ─── Root & Health ──────────────────────────────────────────────────────────── @app.head("/", response_class=HTMLResponse) def root_head(): return HTMLResponse(content="", status_code=200) @app.get("/", response_class=HTMLResponse) def root(): db_status = "Connected" if db_pool is not None else "Disconnected" db_dot_class = "dot-connected" if db_pool is not None else "dot-disconnected" redis_status = "Connected" if redis_conn is not None else "Disconnected" redis_dot_class = "dot-connected" if redis_conn is not None else "dot-disconnected" html_content = f""" AnimalMind Backend Gateway
🐾

AnimalMind Backend

v1.4.0 • Gateway

FastAPI machine learning engine for pet voice classification, breed detection, and posture analysis.

PostgreSQL {db_status}
Redis Cache {redis_status}
Access API Documentation
Running in Hugging Face Spaces Sandbox
""" return HTMLResponse(content=html_content, status_code=200) @app.get("/health") @app.head("/health") def health(): return {"status": "healthy"} if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)