| |
| """SAM 2.1 polygon annotation — GPU run. |
| Detection bbox'larını prompt olarak kullanıp RLE mask üretir. |
| Output: her kaynağın detection.coco.json dosyasına 'segmentation' alanı ekler. |
| """ |
| import json, os, argparse, time |
| from pathlib import Path |
|
|
| ROOT = Path("/arf/scratch/stakan/hitit-proje") |
| SOURCES = ROOT / "datasets" / "sources" |
|
|
| def main(): |
| import torch |
| from PIL import Image |
| import numpy as np |
| |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--source', default='all', help='all|hitit_local|compvis|yeni_veri|deepscribe') |
| ap.add_argument('--limit', type=int, default=0) |
| args = ap.parse_args() |
| |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| print(f"Device: {device}", flush=True) |
| |
| |
| try: |
| from ultralytics import SAM |
| sam = SAM("sam2.1_l.pt") |
| model_name = "sam2.1_l (ultralytics)" |
| except Exception as e: |
| print(f"Ultralytics SAM yüklenmedi: {e}; alternatif deneniyor", flush=True) |
| try: |
| from segment_anything import SamPredictor, sam_model_registry |
| print("segment_anything varsa — manual checkpoint path gerekir") |
| return |
| except ImportError: |
| print("SAM kurulu değil. Install: pip install ultralytics", flush=True) |
| return |
| |
| source_dirs = ['hitit_local', 'compvis', 'yeni_veri', 'deepscribe'] if args.source == 'all' else [args.source] |
| |
| try: |
| from pycocotools import mask as mask_utils |
| except ImportError: |
| os.system("pip install --user --quiet pycocotools") |
| from pycocotools import mask as mask_utils |
| |
| t0 = time.time() |
| total_done = 0 |
| |
| for src_name in source_dirs: |
| src = SOURCES / src_name |
| mp = src / "manifest_detection.jsonl" |
| if not mp.exists(): continue |
| |
| records = [] |
| with open(mp) as f: |
| for line in f: |
| records.append(json.loads(line)) |
| if args.limit: |
| records = records[:args.limit] |
| |
| print(f"\n{src_name}: {len(records)} record", flush=True) |
| |
| out_json = src / "detection.coco.json" |
| coco = json.load(open(out_json)) if out_json.exists() else { |
| "info": {"description": f"{src_name} detection with SAM 2.1 masks"}, |
| "images": [], "annotations": [], "categories": [] |
| } |
| |
| ann_id = 1 + (max([a['id'] for a in coco.get('annotations', [])], default=0)) |
| cat_map = {c['name']: c['id'] for c in coco.get('categories', [])} |
| img_id = 1 + (max([i['id'] for i in coco.get('images', [])], default=0)) |
| existing_img_paths = {i.get('file_name') for i in coco.get('images', [])} |
| |
| for ri, r in enumerate(records): |
| p = r.get('path') |
| if not p or not os.path.exists(p): continue |
| if p in existing_img_paths: continue |
| |
| |
| bboxes_xyxy = [] |
| labels = [] |
| w = r.get('width') or 0 |
| h = r.get('height') or 0 |
| if not (w and h): |
| try: |
| with Image.open(p) as img: |
| w, h = img.size |
| except: continue |
| |
| lp = r.get('label_path') |
| if lp and os.path.exists(lp): |
| try: |
| with open(lp) as f: |
| for ln in f: |
| parts = ln.strip().split() |
| if len(parts) == 5: |
| cls, cx, cy, bw, bh = parts |
| cx, cy, bw, bh = float(cx)*w, float(cy)*h, float(bw)*w, float(bh)*h |
| x1, y1, x2, y2 = cx-bw/2, cy-bh/2, cx+bw/2, cy+bh/2 |
| bboxes_xyxy.append([x1, y1, x2, y2]) |
| labels.append(f"class_{cls}") |
| except: pass |
| extra = r.get('extra') or {} |
| if isinstance(extra, dict): |
| for b in extra.get('bboxes', []): |
| if not isinstance(b, dict): continue |
| lbl = b.get('class_name') or b.get('mzl_label') or b.get('train_label') or 'unknown' |
| yolo = b.get('yolo_bbox') |
| xyxy = b.get('bbox') |
| if yolo and len(yolo) == 4: |
| cx, cy, bw, bh = yolo |
| cx, cy, bw, bh = cx*w, cy*h, bw*w, bh*h |
| bboxes_xyxy.append([cx-bw/2, cy-bh/2, cx+bw/2, cy+bh/2]) |
| labels.append(str(lbl)) |
| elif xyxy and len(xyxy) == 4: |
| bboxes_xyxy.append([float(x) for x in xyxy]) |
| labels.append(str(lbl)) |
| |
| if not bboxes_xyxy: continue |
| |
| |
| try: |
| results = sam(p, bboxes=bboxes_xyxy, verbose=False) |
| except Exception as e: |
| print(f" skip {p}: {e}", flush=True) |
| continue |
| |
| |
| coco['images'].append({ |
| "id": img_id, |
| "file_name": p, |
| "width": w, "height": h, |
| "fold": r.get('fold'), |
| "tablet_id": r.get('tablet_id'), |
| }) |
| |
| |
| if results and hasattr(results[0], 'masks') and results[0].masks is not None: |
| masks = results[0].masks.data.cpu().numpy() |
| for mi, (mask, lbl, bb) in enumerate(zip(masks, labels, bboxes_xyxy)): |
| if lbl not in cat_map: |
| cat_map[lbl] = len(cat_map) + 1 |
| coco['categories'].append({"id": cat_map[lbl], "name": lbl, "supercategory": "sign"}) |
| |
| mask_u8 = (mask > 0).astype(np.uint8) |
| rle = mask_utils.encode(np.asfortranarray(mask_u8)) |
| rle['counts'] = rle['counts'].decode('utf-8') |
| x1, y1, x2, y2 = bb |
| coco['annotations'].append({ |
| "id": ann_id, |
| "image_id": img_id, |
| "category_id": cat_map[lbl], |
| "bbox": [x1, y1, x2-x1, y2-y1], |
| "area": float(mask_u8.sum()), |
| "iscrowd": 0, |
| "segmentation": rle, |
| }) |
| ann_id += 1 |
| img_id += 1 |
| total_done += 1 |
| |
| if (ri+1) % 100 == 0: |
| elapsed = time.time() - t0 |
| rate = total_done / max(elapsed, 1) |
| print(f" [{src_name}] {ri+1}/{len(records)} ({rate:.1f} img/s)", flush=True) |
| |
| with open(out_json, 'w') as f: |
| json.dump(coco, f, ensure_ascii=False) |
| print(f" {src_name}: yazıldı {out_json} — {len(coco['images'])} img, {len(coco['annotations'])} ann", flush=True) |
|
|
| if __name__ == '__main__': |
| main() |
|
|
