Файл web_demo.html отсутствует
" + +@app.get("/web/advanced", response_class=HTMLResponse) +async def web_interface_advanced(): + """Продвинутый веб-интерфейс с Box + Brush промптами""" + web_demo_path = os.path.join(os.path.dirname(__file__), "web_demo_advanced.html") + if os.path.exists(web_demo_path): + with open(web_demo_path, "r", encoding="utf-8") as f: + return f.read() + else: + return "Файл web_demo_advanced.html отсутствует
" + +@app.get("/health") +async def health(): + """Проверка что всё ок""" + return { + "status": "healthy" if predictor is not None else "model not loaded", + "device": str(device) if device else "unknown" + } + +def process_image(image_bytes: bytes) -> np.ndarray: + """Конвертирует байты в numpy array""" + image = Image.open(io.BytesIO(image_bytes)) + if image.mode != "RGB": + image = image.convert("RGB") + return np.array(image) + +def masks_to_coords(masks: np.ndarray, include_contours: bool = False) -> List[Dict[str, Any]]: + """ + Конвертирует маски в координаты bounding box и контуров. + masks: (N, H, W) - N масок + include_contours: если True, добавляет контуры масок + """ + results = [] + + for i, mask in enumerate(masks): + # Находим координаты пикселей маски + y_coords, x_coords = np.where(mask > 0) + + if len(x_coords) == 0: + continue + + # Bounding box + x_min, x_max = int(x_coords.min()), int(x_coords.max()) + y_min, y_max = int(y_coords.min()), int(y_coords.max()) + + # Площадь сегмента + area = int(mask.sum()) + + segment_data = { + "segment_id": i, + "bbox": { + "x_min": x_min, + "y_min": y_min, + "x_max": x_max, + "y_max": y_max, + "width": x_max - x_min, + "height": y_max - y_min + }, + "area": area, + "center": { + "x": int(x_coords.mean()), + "y": int(y_coords.mean()) + } + } + + # Добавляем контуры если нужно + if include_contours: + try: + # Конвертируем маску в uint8 (защита от булевых масок) + if mask.dtype == bool: + mask_uint8 = mask.astype(np.uint8) * 255 + else: + mask_uint8 = (mask * 255).astype(np.uint8) + + # Находим контуры с иерархией для поддержки "дыр" + # RETR_CCOMP: находит внешние контуры И внутренние дыры (holes) + # CHAIN_APPROX_NONE: сохраняет ВСЕ точки для pixel-perfect результата + contours, hierarchy = cv2.findContours(mask_uint8, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE) + except Exception as e: + logger.warning(f"Ошибка при извлечении контуров: {e}, использую fallback") + # Fallback на простое извлечение без иерархии + contours, hierarchy = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) + hierarchy = None + + # Конвертируем контуры в список точек с учетом иерархии + contour_data = [] + + if hierarchy is not None and len(contours) > 0: + hierarchy = hierarchy[0] # OpenCV возвращает hierarchy в странном формате + + for i, contour in enumerate(contours): + try: + # Небольшое упрощение только для очень больших контуров + if len(contour) > 1000: + arc_length = cv2.arcLength(contour, True) + if arc_length > 0: # Защита от деления на 0 + epsilon = 0.0005 * arc_length + approx = cv2.approxPolyDP(contour, epsilon, True) + else: + approx = contour + else: + approx = contour + + # Конвертируем в список [x, y] + points = [[int(point[0][0]), int(point[0][1])] for point in approx] + + if len(points) > 2: + # hierarchy[i] = [Next, Previous, First_Child, Parent] + # Если Parent == -1, это внешний контур + # Если Parent >= 0, это дыра (hole) внутри родительского контура + is_hole = hierarchy[i][3] != -1 + + contour_data.append({ + "points": points, + "is_hole": is_hole + }) + except Exception as e: + logger.warning(f"Ошибка при обработке контура {i}: {e}") + continue + else: + # Fallback если hierarchy не вернулась + for contour in contours: + try: + if len(contour) > 1000: + arc_length = cv2.arcLength(contour, True) + if arc_length > 0: + epsilon = 0.0005 * arc_length + approx = cv2.approxPolyDP(contour, epsilon, True) + else: + approx = contour + else: + approx = contour + + points = [[int(point[0][0]), int(point[0][1])] for point in approx] + if len(points) > 2: + contour_data.append({ + "points": points, + "is_hole": False + }) + except Exception as e: + logger.warning(f"Ошибка при обработке контура: {e}") + continue + + segment_data["contours"] = contour_data if len(contour_data) > 0 else [] + + # Также добавляем RLE (Run-Length Encoding) для компактного представления + # Это полезно если нужно восстановить точную маску + segment_data["mask_rle"] = mask_to_rle(mask) + + results.append(segment_data) + + return results + +def mask_to_rle(mask: np.ndarray) -> Dict[str, Any]: + """ + Конвертирует бинарную маску в RLE (Run-Length Encoding) + Компактное представление маски + """ + # Конвертируем в int если это bool + if mask.dtype == bool: + pixels = mask.astype(np.uint8).flatten() + else: + pixels = mask.flatten() + + pixels = np.concatenate([[0], pixels, [0]]) + runs = np.where(pixels[1:] != pixels[:-1])[0] + 1 + runs[1::2] -= runs[::2] + + return { + "counts": [int(x) for x in runs], # Конвертируем numpy int в Python int + "size": [int(x) for x in mask.shape] # Конвертируем в Python int + } + +def convert_to_native_types(obj): + """ + Рекурсивно конвертирует numpy типы в нативные Python типы + Нужно для сериализации в JSON через FastAPI + """ + if isinstance(obj, np.integer): + return int(obj) + elif isinstance(obj, np.floating): + return float(obj) + elif isinstance(obj, np.ndarray): + return obj.tolist() + elif isinstance(obj, np.bool_): + return bool(obj) + elif isinstance(obj, dict): + return {key: convert_to_native_types(value) for key, value in obj.items()} + elif isinstance(obj, list): + return [convert_to_native_types(item) for item in obj] + return obj + +def clean_mask(mask: np.ndarray, min_area: int = 100) -> np.ndarray: + """ + Очищает маску от мелких артефактов и дыр. + + mask: бинарная маска (H, W) + min_area: минимальная площадь компонента в пикселях + + Returns: очищенная маска + """ + # Конвертируем в uint8 если нужно + if mask.dtype == bool: + mask_uint8 = mask.astype(np.uint8) * 255 + else: + mask_uint8 = (mask * 255).astype(np.uint8) + + # Морфологическое закрытие для удаления мелких дыр + kernel = np.ones((3, 3), np.uint8) + mask_uint8 = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel, iterations=2) + + # Морфологическое открытие для удаления мелких шумов + mask_uint8 = cv2.morphologyEx(mask_uint8, cv2.MORPH_OPEN, kernel, iterations=1) + + # Находим все связанные компоненты + num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask_uint8, connectivity=8) + + # Создаем чистую маску + clean_mask = np.zeros_like(mask_uint8) + + # Оставляем только большие компоненты + for i in range(1, num_labels): # Пропускаем фон (0) + area = stats[i, cv2.CC_STAT_AREA] + if area >= min_area: + clean_mask[labels == i] = 255 + + # Если ничего не осталось, возвращаем самый большой компонент + if clean_mask.sum() == 0 and num_labels > 1: + # Находим самый большой компонент + largest_component = 1 + np.argmax([stats[i, cv2.CC_STAT_AREA] for i in range(1, num_labels)]) + clean_mask[labels == largest_component] = 255 + + return (clean_mask > 127).astype(bool) + +def extract_object_image(image: np.ndarray, mask: np.ndarray, clean: bool = True) -> str: + """ + Вырезает объект из изображения по маске и возвращает base64 PNG с прозрачностью. + + image: RGB изображение (H, W, 3) + mask: бинарная маска (H, W) + clean: применить постобработку для удаления артефактов + + Returns: base64 строка PNG изображения с альфа-каналом + """ + # Конвертируем маску в bool если нужно + if mask.dtype != bool: + mask = mask > 0.5 + + # Очищаем маску от артефактов + if clean: + mask = clean_mask(mask, min_area=100) + + # Создаем RGBA изображение + h, w = image.shape[:2] + rgba = np.zeros((h, w, 4), dtype=np.uint8) + rgba[:, :, :3] = image # RGB каналы + rgba[:, :, 3] = (mask * 255).astype(np.uint8) # Alpha канал из маски + + # Конвертируем в PIL Image + pil_image = Image.fromarray(rgba, 'RGBA') + + # Конвертируем в base64 + buffer = io.BytesIO() + pil_image.save(buffer, format='PNG') + buffer.seek(0) + img_base64 = base64.b64encode(buffer.read()).decode('utf-8') + + return f"data:image/png;base64,{img_base64}" + +@app.post("/segment") +async def segment_image( + file: UploadFile = File(...), + point_x: List[float] = Query(None, description="X координаты точек промпта"), + point_y: List[float] = Query(None, description="Y координаты точек промпта"), + point_labels: List[int] = Query(None, description="Лейблы точек (1=foreground, 0=background)"), + box_x1: float = Query(None, description="X координата левого верхнего угла бокса"), + box_y1: float = Query(None, description="Y координата левого верхнего угла бокса"), + box_x2: float = Query(None, description="X координата правого нижнего угла бокса"), + box_y2: float = Query(None, description="Y координата правого нижнего угла бокса"), + mask_data: str = Query(None, description="Base64 закодированная маска (PNG с альфа-каналом)"), + include_masks: bool = Query(True, description="Включить контуры масок в ответ"), + extract_objects: bool = Query(False, description="Вернуть вырезанные объекты как base64 PNG"), +): + """ + Сегментирует изображение по промпту (точкам, боксу, маске или их комбинации). + + Поддерживаемые промпты: + - Точки (point_x, point_y, point_labels) - клики пользователя + - Бокс (box_x1, box_y1, box_x2, box_y2) - прямоугольное выделение + - Маска (mask_data) - нарисованная кистью маска (зеленый=foreground, красный=background) + - Комбинация промптов - для максимальной точности + + Если промпты не указаны, сегментирует центральный объект. + Если include_masks=True, возвращает контуры масок для точной отрисовки. + Если extract_objects=True, возвращает готовые вырезанные объекты как base64 PNG. + """ + if predictor is None: + raise HTTPException(status_code=503, detail="Модель не загружена, перезапусти сервер") + + try: + # Читаем изображение + image_bytes = await file.read() + image = process_image(image_bytes) + + logger.info(f"Обрабатываю изображение: {image.shape}") + logger.info(f"Параметры: include_masks={include_masks}, extract_objects={extract_objects}") + + # Устанавливаем изображение в предиктор + predictor.set_image(image) + + # Подготавливаем промпты + points = None + labels = None + box = None + + # Проверяем наличие точек + if point_x and point_y: + if len(point_x) != len(point_y): + raise HTTPException(status_code=400, detail="Количество X и Y координат должно совпадать") + points = np.array([[x, y] for x, y in zip(point_x, point_y)]) + labels = np.array(point_labels) if point_labels else np.ones(len(points)) + logger.info(f"Промпт: {len(points)} точек") + + # Проверяем наличие бокса + if all(v is not None for v in [box_x1, box_y1, box_x2, box_y2]): + box = np.array([box_x1, box_y1, box_x2, box_y2]) + logger.info(f"Промпт: бокс [{box_x1:.1f}, {box_y1:.1f}, {box_x2:.1f}, {box_y2:.1f}]") + + # Валидация бокса + if box_x2 <= box_x1 or box_y2 <= box_y1: + raise HTTPException( + status_code=400, + detail="Некорректный бокс: x2 должен быть больше x1, y2 больше y1" + ) + + # Проверяем наличие нарисованной маски + if mask_data: + logger.info("Обрабатываю нарисованную маску...") + try: + # Декодируем base64 + if ',' in mask_data: + mask_data = mask_data.split(',')[1] # Убираем data:image/png;base64, + + mask_bytes = base64.b64decode(mask_data) + mask_image = Image.open(io.BytesIO(mask_bytes)).convert('RGBA') + mask_array = np.array(mask_image) + + # Извлекаем foreground и background пиксели + # Поддерживаем несколько форматов: + # 1. Зеленый (R<100, G>150, B<100) - классический foreground + # 2. Белый/светлый (R>200, G>200, B>200) - часто используется фронтами + # 3. Красный (R>150, G<100, B<100) - background + + green_mask = (mask_array[:, :, 0] < 100) & (mask_array[:, :, 1] > 150) & (mask_array[:, :, 2] < 100) & (mask_array[:, :, 3] > 0) + white_mask = (mask_array[:, :, 0] > 200) & (mask_array[:, :, 1] > 200) & (mask_array[:, :, 2] > 200) & (mask_array[:, :, 3] > 0) + red_mask = (mask_array[:, :, 0] > 150) & (mask_array[:, :, 1] < 100) & (mask_array[:, :, 2] < 100) & (mask_array[:, :, 3] > 0) + + # Объединяем зеленые и белые как foreground + foreground_mask = green_mask | white_mask + + # Сэмплируем точки из закрашенных областей + mask_points = [] + mask_labels = [] + + # Foreground точки (зеленые + белые) + foreground_coords = np.argwhere(foreground_mask) + if len(foreground_coords) > 0: + # Масштабируем к размеру исходного изображения + scale_y = image.shape[0] / mask_array.shape[0] + scale_x = image.shape[1] / mask_array.shape[1] + + # Сэмплируем до 20 точек равномерно (меньше = стабильнее) + step = max(1, len(foreground_coords) // 20) + sampled = foreground_coords[::step][:20] # Максимум 20 точек + + for y, x in sampled: + mask_points.append([x * scale_x, y * scale_y]) + mask_labels.append(1) # foreground + + # Background точки (красные) + red_coords = np.argwhere(red_mask) + if len(red_coords) > 0: + scale_y = image.shape[0] / mask_array.shape[0] + scale_x = image.shape[1] / mask_array.shape[1] + + step = max(1, len(red_coords) // 20) + sampled = red_coords[::step][:20] # Максимум 20 точек + + for y, x in sampled: + mask_points.append([x * scale_x, y * scale_y]) + mask_labels.append(0) # background + + if mask_points: + # Объединяем с существующими точками + if points is not None: + points = np.vstack([points, np.array(mask_points)]) + labels = np.concatenate([labels, np.array(mask_labels)]) + else: + points = np.array(mask_points) + labels = np.array(mask_labels) + + logger.info(f"Промпт из маски: {len(mask_points)} точек ({np.sum(np.array(mask_labels) == 1)} foreground, {np.sum(np.array(mask_labels) == 0)} background)") + else: + logger.warning("Маска пустая или не содержит foreground (зеленых/белых) или background (красных) пикселей") + + except Exception as e: + logger.error(f"Ошибка обработки маски: {e}") + raise HTTPException(status_code=400, detail=f"Некорректная маска: {str(e)}") + + # Делаем предсказание с промптами + if points is not None or box is not None: + logger.info(f"Используем промпты: points={points is not None}, box={box is not None}") + + # Если много точек (>10), используем single mask для стабильности + # Если мало точек или только box, используем multimask для вариативности + use_multimask = True + if points is not None and len(points) > 10: + use_multimask = False + logger.info("Много точек, используем single mask mode для стабильности") + + masks, scores, logits = predictor.predict( + point_coords=points, + point_labels=labels, + box=box, + multimask_output=use_multimask, + ) + + # Если multimask, выбираем лучшую по score + if use_multimask and len(masks) > 1: + best_idx = np.argmax(scores) + masks = masks[best_idx:best_idx+1] + scores = scores[best_idx:best_idx+1] + logger.info(f"Выбрана маска {best_idx} с confidence {scores[0]:.3f}") + else: + # Автоматическая сегментация - берем центральную точку + logger.info("Промпты не указаны, сегментирую центральный объект") + h, w = image.shape[:2] + point = np.array([[w // 2, h // 2]]) + label = np.array([1]) + + masks, scores, logits = predictor.predict( + point_coords=point, + point_labels=label, + multimask_output=True, + ) + + # Конвертируем маски в координаты (с контурами если нужно) + segments = masks_to_coords(masks, include_contours=include_masks) + + logger.info(f"Найдено сегментов: {len(segments)}, масок: {len(masks)}") + logger.info(f"extract_objects = {extract_objects}") + + # Добавляем confidence scores + for i, seg in enumerate(segments): + seg["confidence"] = float(scores[i]) if i < len(scores) else 0.0 + + # Если нужно - вырезаем объект и добавляем base64 + logger.info(f"Обрабатываю сегмент {i}: extract_objects={extract_objects}, i < len(masks) = {i < len(masks)}") + if extract_objects and i < len(masks): + logger.info(f"Вырезаю объект {i}...") + seg["extracted_image"] = extract_object_image(image, masks[i]) + logger.info(f"✓ Вырезан объект {i}, размер маски: {masks[i].sum()} пикселей") + else: + logger.warning(f"❌ Пропускаю объект {i}: extract_objects={extract_objects}") + + result = { + "success": True, + "image_size": { + "width": int(image.shape[1]), + "height": int(image.shape[0]) + }, + "segments_count": len(segments), + "segments": segments + } + + # Конвертируем все numpy типы в нативные Python типы + return convert_to_native_types(result) + + except Exception as e: + logger.error(f"Ошибка при сегментации: {e}") + raise HTTPException(status_code=500, detail=f"Ошибка обработки: {str(e)}") + +@app.post("/segment/auto") +async def segment_auto( + file: UploadFile = File(...), + points_per_side: int = Query(32, description="Количество точек на сторону для автосегментации"), + include_masks: bool = Query(True, description="Включить контуры масок в ответ"), +): + """ + Автоматическая сегментация всех объектов на изображении. + Использует grid of points для поиска всех возможных объектов. + Если include_masks=True, возвращает контуры масок для точной отрисовки. + """ + if predictor is None: + raise HTTPException(status_code=503, detail="Модель не загружена") + + try: + image_bytes = await file.read() + image = process_image(image_bytes) + + logger.info(f"Автосегментация изображения: {image.shape}") + + predictor.set_image(image) + + # Создаем сетку точек + h, w = image.shape[:2] + x_coords = np.linspace(0, w, points_per_side) + y_coords = np.linspace(0, h, points_per_side) + + all_segments = [] + segment_id = 0 + + # Для каждой точки в сетке пытаемся найти объект + for y in y_coords: + for x in x_coords: + point = np.array([[x, y]]) + label = np.array([1]) + + masks, scores, _ = predictor.predict( + point_coords=point, + point_labels=label, + multimask_output=False, + ) + + if masks.shape[0] > 0 and scores[0] > 0.5: # Порог confidence + segments = masks_to_coords(masks, include_contours=include_masks) + for seg in segments: + seg["segment_id"] = segment_id + seg["confidence"] = float(scores[0]) + all_segments.append(seg) + segment_id += 1 + + # Убираем дубликаты (примерно) + # Два сегмента считаем дубликатами если их центры близко + unique_segments = [] + for seg in all_segments: + is_duplicate = False + for unique_seg in unique_segments: + dx = seg["center"]["x"] - unique_seg["center"]["x"] + dy = seg["center"]["y"] - unique_seg["center"]["y"] + dist = (dx**2 + dy**2) ** 0.5 + + if dist < 50: # Порог расстояния между центрами + is_duplicate = True + break + + if not is_duplicate: + unique_segments.append(seg) + + result = { + "success": True, + "image_size": { + "width": int(image.shape[1]), + "height": int(image.shape[0]) + }, + "segments_count": len(unique_segments), + "segments": unique_segments + } + + # Конвертируем все numpy типы в нативные Python типы + return convert_to_native_types(result) + + except Exception as e: + logger.error(f"Ошибка при автосегментации: {e}") + raise HTTPException(status_code=500, detail=f"Ошибка обработки: {str(e)}") + +@app.post("/segment/batch", response_model=BatchSegmentResponse) +async def segment_batch(request: BatchSegmentRequest = Body(...)): + """ + Батчинг сегментация нескольких объектов. + + Принимает изображение и массив промптов (mask/box/points). + Обрабатывает каждый selected промпт отдельно. + Возвращает массив результатов с метаданными. + + Идеально для: + - Множественных объектов + - Мобильных приложений + - Когда фронт уже разделил объекты + """ + if predictor is None: + raise HTTPException(status_code=503, detail="Модель не загружена, перезапусти сервер") + + try: + # Декодируем изображение из base64 + image_data = request.image + if ',' in image_data: + image_data = image_data.split(',')[1] # Убираем data:image/...;base64, + + image_bytes = base64.b64decode(image_data) + image = process_image(image_bytes) + + logger.info(f"Батчинг сегментация: {image.shape}, промптов: {len(request.prompts)}") + + # Устанавливаем изображение один раз + predictor.set_image(image) + + results = [] + + # Фильтруем только selected промпты + selected_prompts = [p for p in request.prompts if p.selected] + logger.info(f"Обрабатываем {len(selected_prompts)} из {len(request.prompts)} промптов") + + # Обрабатываем каждый промпт отдельно + for prompt in selected_prompts: + logger.info(f"Обрабатываю промпт #{prompt.id}, тип: {prompt.type}, label: {prompt.label}") + + try: + # Подготавливаем промпт в зависимости от типа + points = None + labels = None + box = None + + if prompt.type == "mask": + # Декодируем маску и извлекаем точки + mask_data = prompt.data + if ',' in mask_data: + mask_data = mask_data.split(',')[1] + + mask_bytes = base64.b64decode(mask_data) + mask_image = Image.open(io.BytesIO(mask_bytes)).convert('RGBA') + mask_array = np.array(mask_image) + + # Извлекаем foreground и background пиксели + # Поддерживаем несколько форматов: + # 1. Зеленый (R<100, G>150, B<100) - классический foreground + # 2. Белый/светлый (R>200, G>200, B>200) - часто используется фронтами + # 3. Красный (R>150, G<100, B<100) - background + + green_mask = (mask_array[:, :, 0] < 100) & (mask_array[:, :, 1] > 150) & (mask_array[:, :, 2] < 100) & (mask_array[:, :, 3] > 0) + white_mask = (mask_array[:, :, 0] > 200) & (mask_array[:, :, 1] > 200) & (mask_array[:, :, 2] > 200) & (mask_array[:, :, 3] > 0) + red_mask = (mask_array[:, :, 0] > 150) & (mask_array[:, :, 1] < 100) & (mask_array[:, :, 2] < 100) & (mask_array[:, :, 3] > 0) + + # Объединяем зеленые и белые как foreground + foreground_mask = green_mask | white_mask + + mask_points = [] + mask_labels = [] + + # Foreground точки (зеленые + белые) + foreground_coords = np.argwhere(foreground_mask) + if len(foreground_coords) > 0: + scale_y = image.shape[0] / mask_array.shape[0] + scale_x = image.shape[1] / mask_array.shape[1] + step = max(1, len(foreground_coords) // 20) + sampled = foreground_coords[::step][:20] + + for y, x in sampled: + mask_points.append([x * scale_x, y * scale_y]) + mask_labels.append(1) + + # Background точки + red_coords = np.argwhere(red_mask) + if len(red_coords) > 0: + scale_y = image.shape[0] / mask_array.shape[0] + scale_x = image.shape[1] / mask_array.shape[1] + step = max(1, len(red_coords) // 20) + sampled = red_coords[::step][:20] + + for y, x in sampled: + mask_points.append([x * scale_x, y * scale_y]) + mask_labels.append(0) + + if mask_points: + points = np.array(mask_points) + labels = np.array(mask_labels) + + elif prompt.type == "box": + # Парсим bbox - может быть нормализованный (0-1) или пиксельный + bbox_data = prompt.bbox if prompt.bbox else None + + if bbox_data: + x1 = bbox_data.x_min + y1 = bbox_data.y_min + x2 = bbox_data.x_max + y2 = bbox_data.y_max + + # Если нормализованные координаты (0-1), конвертируем в пиксели + if x2 <= 1.0 and y2 <= 1.0: + x1 *= image.shape[1] + x2 *= image.shape[1] + y1 *= image.shape[0] + y2 *= image.shape[0] + + box = np.array([x1, y1, x2, y2]) + + elif prompt.type == "points": + # Ожидаем JSON в формате [[x, y, label], ...] + import json + points_data = json.loads(prompt.data) + + points_list = [] + labels_list = [] + + for point in points_data: + x, y = point[0], point[1] + label = point[2] if len(point) > 2 else 1 + + # Если нормализованные, конвертируем + if x <= 1.0 and y <= 1.0: + x *= image.shape[1] + y *= image.shape[0] + + points_list.append([x, y]) + labels_list.append(label) + + points = np.array(points_list) + labels = np.array(labels_list) + + # Делаем предсказание + if points is not None or box is not None: + # Решаем использовать ли multimask + use_multimask = True + if points is not None and len(points) > 10: + use_multimask = False + + masks, scores, logits = predictor.predict( + point_coords=points, + point_labels=labels, + box=box, + multimask_output=use_multimask, + ) + + # Если multimask, выбираем лучшую + if use_multimask and len(masks) > 1: + best_idx = np.argmax(scores) + masks = masks[best_idx:best_idx+1] + scores = scores[best_idx:best_idx+1] + + # Берем первую маску + mask = masks[0] + score = float(scores[0]) + + # Очищаем маску если нужно + if request.options.clean_masks: + mask = clean_mask(mask, min_area=100) + + # Вычисляем метрики + y_coords, x_coords = np.where(mask > 0) + + if len(x_coords) > 0: + x_min, x_max = int(x_coords.min()), int(x_coords.max()) + y_min, y_max = int(y_coords.min()), int(y_coords.max()) + area = int(mask.sum()) + center_x = int(x_coords.mean()) + center_y = int(y_coords.mean()) + + # Формируем результат + result = { + "id": prompt.id, + "label": prompt.label, + "bbox": { + "x_min": x_min, + "y_min": y_min, + "x_max": x_max, + "y_max": y_max, + "width": x_max - x_min, + "height": y_max - y_min + }, + "area": area, + "center": { + "x": center_x, + "y": center_y + }, + "confidence": score + } + + # Добавляем вырезанный объект если нужно + if request.options.extract_objects: + result["extracted_image"] = extract_object_image( + image, mask, clean=request.options.clean_masks + ) + + # Добавляем контуры если нужно + if request.options.include_masks: + segments = masks_to_coords(masks, include_contours=True) + if segments: + result["contours"] = segments[0].get("contours", []) + result["mask_rle"] = segments[0].get("mask_rle", {}) + + results.append(result) + logger.info(f"✓ Промпт #{prompt.id} обработан, confidence: {score:.3f}") + else: + logger.warning(f"✗ Промпт #{prompt.id} не дал результата") + else: + logger.warning(f"✗ Промпт #{prompt.id}: нет данных для сегментации") + + except Exception as e: + logger.error(f"✗ Ошибка обработки промпта #{prompt.id}: {e}") + # Продолжаем обработку остальных промптов + continue + + response = { + "success": True, + "image_size": { + "width": int(image.shape[1]), + "height": int(image.shape[0]) + }, + "results": results + } + + logger.info(f"Батчинг завершен: обработано {len(results)} объектов") + + # Сохраняем лог запроса для аудита (только метаданные, без изображений) + try: + request_dict = request.dict() + save_batch_request_log(request_dict, response, image.shape[1], image.shape[0]) + except Exception as e: + logger.warning(f"Не удалось сохранить лог запроса: {e}") + + return convert_to_native_types(response) + + except Exception as e: + logger.error(f"Ошибка при батчинг сегментации: {e}") + raise HTTPException(status_code=500, detail=f"Ошибка обработки: {str(e)}") + +if __name__ == "__main__": + import uvicorn + import os + + # Порт из переменной окружения (для HF Spaces) или 8000 по умолчанию + port = int(os.getenv("PORT", 8000)) + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/download_model.py b/download_model.py new file mode 100755 index 0000000000000000000000000000000000000000..82fce62e0a161b7a70ac6989840c43b47484d75a --- /dev/null +++ b/download_model.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +""" +Скрипт для скачивания модели SAM2. +Блин, Facebook не может нормально в pip packaging, поэтому качаем руками. +""" + +import os +import urllib.request +import sys + +# Директория для чекпоинтов +CHECKPOINT_DIR = "checkpoints" +os.makedirs(CHECKPOINT_DIR, exist_ok=True) + +# Модели на выбор +MODELS = { + "tiny": { + "url": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_tiny.pt", + "filename": "sam2.1_hiera_tiny.pt", + "size": "~39MB" + }, + "small": { + "url": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt", + "filename": "sam2.1_hiera_small.pt", + "size": "~46MB" + }, + "base_plus": { + "url": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_base_plus.pt", + "filename": "sam2.1_hiera_base_plus.pt", + "size": "~81MB" + }, + "large": { + "url": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt", + "filename": "sam2.1_hiera_large.pt", + "size": "~224MB" + } +} + +def download_model(model_name="tiny"): + """Качает модель, показывает прогресс""" + if model_name not in MODELS: + print(f"Неизвестная модель: {model_name}") + print(f"Доступные: {', '.join(MODELS.keys())}") + sys.exit(1) + + model_info = MODELS[model_name] + filepath = os.path.join(CHECKPOINT_DIR, model_info["filename"]) + + if os.path.exists(filepath): + print(f"Модель уже скачана: {filepath}") + return filepath + + print(f"Качаю {model_name} модель ({model_info['size']})...") + print(f"URL: {model_info['url']}") + + def progress_hook(block_num, block_size, total_size): + downloaded = block_num * block_size + if total_size > 0: + percent = min(100, downloaded * 100 / total_size) + sys.stdout.write(f"\rПрогресс: {percent:.1f}%") + sys.stdout.flush() + + try: + urllib.request.urlretrieve( + model_info["url"], + filepath, + reporthook=progress_hook + ) + print(f"\n✓ Модель скачана: {filepath}") + return filepath + except Exception as e: + print(f"\n✗ Ошибка при скачивании: {e}") + if os.path.exists(filepath): + os.remove(filepath) + sys.exit(1) + +if __name__ == "__main__": + model_name = sys.argv[1] if len(sys.argv) > 1 else "tiny" + download_model(model_name) + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..42b894ca10faaa31c21a48a93d7653f94dba16a2 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,10 @@ +fastapi==0.115.0 +uvicorn[standard]==0.32.0 +python-multipart==0.0.12 +Pillow==11.0.0 +numpy==2.1.0 +torch==2.6.0 +torchvision==0.21.0 +opencv-python==4.10.0.84 +pydantic==2.9.0 + diff --git a/sam2_repo/README.md b/sam2_repo/README.md new file mode 100644 index 0000000000000000000000000000000000000000..85a7eb958bced5495ff990c2bcbe7d99662c660f --- /dev/null +++ b/sam2_repo/README.md @@ -0,0 +1,224 @@ +# SAM 2: Segment Anything in Images and Videos + +**[AI at Meta, FAIR](https://ai.meta.com/research/)** + +[Nikhila Ravi](https://nikhilaravi.com/), [Valentin Gabeur](https://gabeur.github.io/), [Yuan-Ting Hu](https://scholar.google.com/citations?user=E8DVVYQAAAAJ&hl=en), [Ronghang Hu](https://ronghanghu.com/), [Chaitanya Ryali](https://scholar.google.com/citations?user=4LWx24UAAAAJ&hl=en), [Tengyu Ma](https://scholar.google.com/citations?user=VeTSl0wAAAAJ&hl=en), [Haitham Khedr](https://hkhedr.com/), [Roman Rädle](https://scholar.google.de/citations?user=Tpt57v0AAAAJ&hl=en), [Chloe Rolland](https://scholar.google.com/citations?hl=fr&user=n-SnMhoAAAAJ), [Laura Gustafson](https://scholar.google.com/citations?user=c8IpF9gAAAAJ&hl=en), [Eric Mintun](https://ericmintun.github.io/), [Junting Pan](https://junting.github.io/), [Kalyan Vasudev Alwala](https://scholar.google.co.in/citations?user=m34oaWEAAAAJ&hl=en), [Nicolas Carion](https://www.nicolascarion.com/), [Chao-Yuan Wu](https://chaoyuan.org/), [Ross Girshick](https://www.rossgirshick.info/), [Piotr Dollár](https://pdollar.github.io/), [Christoph Feichtenhofer](https://feichtenhofer.github.io/) + +[[`Paper`](https://ai.meta.com/research/publications/sam-2-segment-anything-in-images-and-videos/)] [[`Project`](https://ai.meta.com/sam2)] [[`Demo`](https://sam2.metademolab.com/)] [[`Dataset`](https://ai.meta.com/datasets/segment-anything-video)] [[`Blog`](https://ai.meta.com/blog/segment-anything-2)] [[`BibTeX`](#citing-sam-2)] + + + +**Segment Anything Model 2 (SAM 2)** is a foundation model towards solving promptable visual segmentation in images and videos. We extend SAM to video by considering images as a video with a single frame. The model design is a simple transformer architecture with streaming memory for real-time video processing. We build a model-in-the-loop data engine, which improves model and data via user interaction, to collect [**our SA-V dataset**](https://ai.meta.com/datasets/segment-anything-video), the largest video segmentation dataset to date. SAM 2 trained on our data provides strong performance across a wide range of tasks and visual domains. + + + +## Latest updates + +**12/11/2024 -- full model compilation for a major VOS speedup and a new `SAM2VideoPredictor` to better handle multi-object tracking** + +- We now support `torch.compile` of the entire SAM 2 model on videos, which can be turned on by setting `vos_optimized=True` in `build_sam2_video_predictor`, leading to a major speedup for VOS inference. +- We update the implementation of `SAM2VideoPredictor` to support independent per-object inference, allowing us to relax the assumption of prompting for multi-object tracking and adding new objects after tracking starts. +- See [`RELEASE_NOTES.md`](RELEASE_NOTES.md) for full details. + +**09/30/2024 -- SAM 2.1 Developer Suite (new checkpoints, training code, web demo) is released** + +- A new suite of improved model checkpoints (denoted as **SAM 2.1**) are released. See [Model Description](#model-description) for details. + * To use the new SAM 2.1 checkpoints, you need the latest model code from this repo. If you have installed an earlier version of this repo, please first uninstall the previous version via `pip uninstall SAM-2`, pull the latest code from this repo (with `git pull`), and then reinstall the repo following [Installation](#installation) below. +- The training (and fine-tuning) code has been released. See [`training/README.md`](training/README.md) on how to get started. +- The frontend + backend code for the SAM 2 web demo has been released. See [`demo/README.md`](demo/README.md) for details. + +## Installation + +SAM 2 needs to be installed first before use. The code requires `python>=3.10`, as well as `torch>=2.5.1` and `torchvision>=0.20.1`. Please follow the instructions [here](https://pytorch.org/get-started/locally/) to install both PyTorch and TorchVision dependencies. You can install SAM 2 on a GPU machine using: + +```bash +git clone https://github.com/facebookresearch/sam2.git && cd sam2 + +pip install -e . +``` +If you are installing on Windows, it's strongly recommended to use [Windows Subsystem for Linux (WSL)](https://learn.microsoft.com/en-us/windows/wsl/install) with Ubuntu. + +To use the SAM 2 predictor and run the example notebooks, `jupyter` and `matplotlib` are required and can be installed by: + +```bash +pip install -e ".[notebooks]" +``` + +Note: +1. It's recommended to create a new Python environment via [Anaconda](https://www.anaconda.com/) for this installation and install PyTorch 2.5.1 (or higher) via `pip` following https://pytorch.org/. If you have a PyTorch version lower than 2.5.1 in your current environment, the installation command above will try to upgrade it to the latest PyTorch version using `pip`. +2. The step above requires compiling a custom CUDA kernel with the `nvcc` compiler. If it isn't already available on your machine, please install the [CUDA toolkits](https://developer.nvidia.com/cuda-toolkit-archive) with a version that matches your PyTorch CUDA version. +3. If you see a message like `Failed to build the SAM 2 CUDA extension` during installation, you can ignore it and still use SAM 2 (some post-processing functionality may be limited, but it doesn't affect the results in most cases). + +Please see [`INSTALL.md`](./INSTALL.md) for FAQs on potential issues and solutions. + +## Getting Started + +### Download Checkpoints + +First, we need to download a model checkpoint. All the model checkpoints can be downloaded by running: + +```bash +cd checkpoints && \ +./download_ckpts.sh && \ +cd .. +``` + +or individually from: + +- [sam2.1_hiera_tiny.pt](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_tiny.pt) +- [sam2.1_hiera_small.pt](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt) +- [sam2.1_hiera_base_plus.pt](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_base_plus.pt) +- [sam2.1_hiera_large.pt](https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt) + +(note that these are the improved checkpoints denoted as SAM 2.1; see [Model Description](#model-description) for details.) + +Then SAM 2 can be used in a few lines as follows for image and video prediction. + +### Image prediction + +SAM 2 has all the capabilities of [SAM](https://github.com/facebookresearch/segment-anything) on static images, and we provide image prediction APIs that closely resemble SAM for image use cases. The `SAM2ImagePredictor` class has an easy interface for image prompting. + +```python +import torch +from sam2.build_sam import build_sam2 +from sam2.sam2_image_predictor import SAM2ImagePredictor + +checkpoint = "./checkpoints/sam2.1_hiera_large.pt" +model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml" +predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint)) + +with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): + predictor.set_image(Выдели объект прямоугольником → Получи точную сегментацию
+ +Кликни или перетащи фото сюда
+ + +Обрабатываю изображение...
+Box выделение ИЛИ рисование кистью (зеленый = объект, красный = фон)
+ +Кликни или перетащи фото сюда
+ + +Обрабатываю изображение...
+