File size: 2,564 Bytes
c881b77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | import os
import json
import cv2
import torch
from torchvision import transforms
from PIL import Image
from collections import defaultdict
import xml.etree.ElementTree as ET
from collections import Counter
PROJECT_DIR = os.getenv('DSP_PROJECT_DIR', '/path/to/DSP_PROJECT_DIR') # Set this manually if the environment variable is unavailable
image_dir = '/path/to/RUOD/RUOD_pic/train' # Replace with your actual path
label_dir = '/path/to/RUOD/RUOD_ANN' # Replace with your actual path
output_dir = os.path.join(PROJECT_DIR, "data", "RUOD", "patches")
os.makedirs(output_dir, exist_ok=True)
annos = json.load(open(os.path.join(label_dir, "instances_train.json"), "r"))
images_items = annos["images"]
annos_items = annos["annotations"]
cates_items = annos["categories"]
catemap = {}
for cate in cates_items:
catemap[cate["id"]] = cate["name"]
files = [i["file_name"] for i in images_items]
labels = defaultdict(list)
for item in annos_items:
image_id = item["image_id"]
filename = files[image_id-1]
labels[filename].append([catemap[item["category_id"]]] + item["bbox"])
print(len(files))
counter = Counter()
for image_name in files:
if not image_name.endswith(".jpg"):
continue
image_path = os.path.join(image_dir, image_name)
image = cv2.imread(image_path)
image_height, image_width, _ = image.shape
lines = labels[image_name]
# if image_name == '008431.jpg':
# import pdb; pdb.set_trace()
for i,line in enumerate(lines):
parts = line
class_name = parts[0]
xmin, ymin, w, h = parts[1:]
bbox_width = w
bbox_height = h
xmax = xmin + w
ymax = ymin + h
bbox_area = bbox_width * bbox_height
image_area = image_width * image_height
bbox_ratio = bbox_area / image_area
if bbox_ratio < 0.001:
continue
class_dir = os.path.join(output_dir, class_name)
os.makedirs(class_dir, exist_ok=True)
counter[class_dir] += 1
xmin, ymin, xmax, ymax = int(xmin), int(ymin), int(xmax), int(ymax)
cropped_image = image[ymin:ymax, xmin:xmax]
output_image_name = f"{image_name[:-4]}_{i}.jpg"
output_image_path = os.path.join(class_dir, output_image_name)
try:
cv2.imwrite(output_image_path, cropped_image)
except:
import pdb; pdb.set_trace()
print(bbox_ratio, xmin, ymin, xmax, ymax, image_name)
print(counter)
|