Spaces:
Sleeping
Sleeping
Hikity
feat: implement currency classification pipeline with EfficientNet-B0 and model training utilities
2ba1249 | import numpy as np | |
| import torch | |
| import cv2 | |
| from torchvision.transforms import Normalize | |
| try: | |
| from albumentations.augmentations.functional import image_compression | |
| except ImportError: | |
| try: | |
| from albumentations.functional import image_compression | |
| except ImportError: | |
| def image_compression(img, *args, **kwargs): | |
| return img | |
| mean = [0.485, 0.456, 0.406] | |
| std = [0.229, 0.224, 0.225] | |
| normalize_transform = Normalize(mean, std) | |
| def confident_strategy(pred, t=0.8): | |
| pred = np.array(pred) | |
| sz = len(pred) | |
| fakes = np.count_nonzero(pred > t) | |
| # If >40% of frames are fake and we have at least 11 fake frames | |
| if fakes > sz // 2.5 and fakes > 11: | |
| return np.mean(pred[pred > t]) | |
| elif np.count_nonzero(pred < 0.2) > 0.9 * sz: | |
| return np.mean(pred[pred < 0.2]) | |
| else: | |
| return np.mean(pred) | |
| def put_to_center(img, input_size): | |
| img = img[:input_size, :input_size] | |
| image = np.zeros((input_size, input_size, 3), dtype=np.uint8) | |
| start_w = (input_size - img.shape[1]) // 2 | |
| start_h = (input_size - img.shape[0]) // 2 | |
| image[start_h:start_h + img.shape[0], start_w: start_w + img.shape[1], :] = img | |
| return image | |
| def isotropically_resize_image(img, size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC): | |
| h, w = img.shape[:2] | |
| if max(w, h) == size: | |
| return img | |
| if w > h: | |
| scale = size / w | |
| h = h * scale | |
| w = size | |
| else: | |
| scale = size / h | |
| w = w * scale | |
| h = size | |
| interpolation = interpolation_up if scale > 1 else interpolation_down | |
| resized = cv2.resize(img, (int(w), int(h)), interpolation=interpolation) | |
| return resized | |
| import os | |
| # Inference chunk size: how many face crops go through the B7 ensemble at | |
| # once. Chunking changes peak memory only, never the prediction (same faces, | |
| # same averaging strategy). 8 keeps a dual-B7 pass under ~2GB on CPU; raise it | |
| # via env on machines with headroom. | |
| INFER_CHUNK = int(os.environ.get("INFER_BATCH_SIZE", "8")) | |
| def predict_on_video(face_extractor, video_path, batch_size, input_size, models, strategy=confident_strategy, | |
| apply_compression=False, device=None): | |
| if device is None: | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| max_faces = batch_size * 4 | |
| try: | |
| faces = face_extractor.process_video(video_path) | |
| faces_detected = 0 | |
| frames_analyzed = len(faces) | |
| if len(faces) > 0: | |
| crops = [] | |
| for frame_data in faces: | |
| faces_detected += len(frame_data["faces"]) | |
| for face in frame_data["faces"]: | |
| if len(crops) >= max_faces: | |
| continue | |
| resized_face = isotropically_resize_image(face, input_size) | |
| resized_face = put_to_center(resized_face, input_size) | |
| if apply_compression: | |
| resized_face = image_compression(resized_face, quality=90, image_type=".jpg") | |
| crops.append(resized_face) | |
| n = len(crops) | |
| if n > 0: | |
| x = np.stack(crops) # only actual faces — no fixed 128-slot buffer | |
| with torch.no_grad(): | |
| preds = [] | |
| for model in models: | |
| frame_preds = [] | |
| # Chunked inference: identical output to a single big | |
| # batch, but peak memory stays bounded on 8GB machines | |
| # and free-tier CPU hosts. | |
| for start in range(0, n, INFER_CHUNK): | |
| chunk = torch.tensor(x[start:start + INFER_CHUNK], device=device).float() | |
| chunk = chunk.permute((0, 3, 1, 2)) | |
| for i in range(len(chunk)): | |
| chunk[i] = normalize_transform(chunk[i] / 255.) | |
| model_input = chunk.half() if device == "cuda" else chunk.float() | |
| y_pred = torch.sigmoid(model(model_input)) | |
| frame_preds.extend(y_pred.view(-1).cpu().numpy().tolist()) | |
| del chunk, model_input, y_pred | |
| preds.append(strategy(np.array(frame_preds))) | |
| return float(np.mean(preds)), faces_detected, frames_analyzed | |
| except Exception as e: | |
| print("Prediction error on video %s: %s" % (video_path, str(e))) | |
| return 0.5, 0, 0 | |