File size: 4,276 Bytes
b5f33fd |
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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import os
from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Tuple
import pandas as pd
IMAGES_EXT: Tuple = (".jpeg", ".jpg", ".png", ".webp", ".bmp", ".gif")
VIDEO_EXT: Tuple = (".mp4", ".avi", ".mov", ".mkv", ".webm")
@dataclass
class PictureInfo:
image_path: str
age: Optional[str] # age or age range(start;end format) or "-1"
gender: Optional[str] # "M" of "F" or "-1"
bbox: List[int] = field(default_factory=lambda: [-1, -1, -1, -1]) # face bbox: xyxy
person_bbox: List[int] = field(default_factory=lambda: [-1, -1, -1, -1]) # person bbox: xyxy
@property
def has_person_bbox(self) -> bool:
return any(coord != -1 for coord in self.person_bbox)
@property
def has_face_bbox(self) -> bool:
return any(coord != -1 for coord in self.bbox)
def has_gt(self, only_age: bool = False) -> bool:
if only_age:
return self.age != "-1"
else:
return not (self.age == "-1" and self.gender == "-1")
def clear_person_bbox(self):
self.person_bbox = [-1, -1, -1, -1]
def clear_face_bbox(self):
self.bbox = [-1, -1, -1, -1]
class AnnotType(Enum):
ORIGINAL = "original"
PERSONS = "persons"
NONE = "none"
@classmethod
def _missing_(cls, value):
print(f"WARN: Unknown annotation type {value}.")
return AnnotType.NONE
def get_all_files(path: str, extensions: Tuple = IMAGES_EXT):
files_all = []
for root, subFolders, files in os.walk(path):
for name in files:
# linux tricks with .directory that still is file
if "directory" not in name and sum([ext.lower() in name.lower() for ext in extensions]) > 0:
files_all.append(os.path.join(root, name))
return files_all
class InputType(Enum):
Image = 0
Video = 1
VideoStream = 2
def get_input_type(input_path: str) -> InputType:
if os.path.isdir(input_path):
print("Input is a folder, only images will be processed")
return InputType.Image
elif os.path.isfile(input_path):
if input_path.endswith(VIDEO_EXT):
return InputType.Video
if input_path.endswith(IMAGES_EXT):
return InputType.Image
else:
raise ValueError(
f"Unknown or unsupported input file format {input_path}, \
supported video formats: {VIDEO_EXT}, \
supported image formats: {IMAGES_EXT}"
)
elif input_path.startswith("http") and not input_path.endswith(IMAGES_EXT):
return InputType.VideoStream
else:
raise ValueError(f"Unknown input {input_path}")
def read_csv_annotation_file(annotation_file: str, images_dir: str, ignore_without_gt=False):
bboxes_per_image: Dict[str, List[PictureInfo]] = defaultdict(list)
df = pd.read_csv(annotation_file, sep=",")
annot_type = AnnotType("persons") if "person_x0" in df.columns else AnnotType("original")
print(f"Reading {annotation_file} (type: {annot_type})...")
missing_images = 0
for index, row in df.iterrows():
img_path = os.path.join(images_dir, row["img_name"])
if not os.path.exists(img_path):
missing_images += 1
continue
face_x1, face_y1, face_x2, face_y2 = row["face_x0"], row["face_y0"], row["face_x1"], row["face_y1"]
age, gender = str(row["age"]), str(row["gender"])
if ignore_without_gt and (age == "-1" or gender == "-1"):
continue
if annot_type == AnnotType.PERSONS:
p_x1, p_y1, p_x2, p_y2 = row["person_x0"], row["person_y0"], row["person_x1"], row["person_y1"]
person_bbox = list(map(int, [p_x1, p_y1, p_x2, p_y2]))
else:
person_bbox = [-1, -1, -1, -1]
bbox = list(map(int, [face_x1, face_y1, face_x2, face_y2]))
pic_info = PictureInfo(img_path, age, gender, bbox, person_bbox)
assert isinstance(pic_info.person_bbox, list)
bboxes_per_image[img_path].append(pic_info)
if missing_images > 0:
print(f"WARNING: Missing images: {missing_images}/{len(df)}")
return bboxes_per_image, annot_type
|