Skip to content

Commit

Permalink
deploy: f7103fe
Browse files Browse the repository at this point in the history
  • Loading branch information
@drcege
drcege committed Aug 22, 2024
1 parent 230ef5d commit fec586d
Show file tree
Hide file tree
Showing 6 changed files with 117 additions and 75 deletions.
51 changes: 31 additions & 20 deletions _modules/data_juicer/ops/filter/image_face_ratio_filter.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,33 +82,43 @@
<div itemprop="articleBody">

<h1>Source code for data_juicer.ops.filter.image_face_ratio_filter</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span></span><span class="kn">import</span> <span class="nn">os</span>

<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">from</span> <span class="nn">jsonargparse.typing</span> <span class="kn">import</span> <span class="n">ClosedUnitInterval</span>
<span class="kn">from</span> <span class="nn">loguru</span> <span class="kn">import</span> <span class="n">logger</span>

<span class="kn">from</span> <span class="nn">data_juicer.utils.availability_utils</span> <span class="kn">import</span> <span class="n">AvailabilityChecking</span>
<span class="kn">from</span> <span class="nn">data_juicer.utils.constant</span> <span class="kn">import</span> <span class="n">Fields</span><span class="p">,</span> <span class="n">StatsKeys</span>
<span class="kn">from</span> <span class="nn">data_juicer.utils.mm_utils</span> <span class="kn">import</span> <span class="p">(</span><span class="n">load_data_with_context</span><span class="p">,</span> <span class="n">load_image</span><span class="p">,</span>
<span class="n">pil_to_opencv</span><span class="p">)</span>
<span class="kn">from</span> <span class="nn">data_juicer.utils.mm_utils</span> <span class="kn">import</span> <span class="p">(</span><span class="n">detect_faces</span><span class="p">,</span> <span class="n">load_data_with_context</span><span class="p">,</span>
<span class="n">load_image</span><span class="p">)</span>
<span class="kn">from</span> <span class="nn">data_juicer.utils.model_utils</span> <span class="kn">import</span> <span class="n">get_model</span><span class="p">,</span> <span class="n">prepare_model</span>

<span class="kn">from</span> <span class="nn">..base_op</span> <span class="kn">import</span> <span class="n">OPERATORS</span><span class="p">,</span> <span class="n">Filter</span>
<span class="kn">from</span> <span class="nn">..base_op</span> <span class="kn">import</span> <span class="n">OPERATORS</span><span class="p">,</span> <span class="n">UNFORKABLE</span><span class="p">,</span> <span class="n">Filter</span>
<span class="kn">from</span> <span class="nn">..op_fusion</span> <span class="kn">import</span> <span class="n">LOADED_IMAGES</span>

<span class="n">OP_NAME</span> <span class="o">=</span> <span class="s1">&#39;image_face_ratio_filter&#39;</span>

<span class="k">with</span> <span class="n">AvailabilityChecking</span><span class="p">([</span><span class="s1">&#39;dlib&#39;</span><span class="p">],</span> <span class="n">OP_NAME</span><span class="p">):</span>
<span class="kn">import</span> <span class="nn">dlib</span>
<span class="k">with</span> <span class="n">AvailabilityChecking</span><span class="p">([</span><span class="s1">&#39;opencv-python&#39;</span><span class="p">],</span> <span class="n">OP_NAME</span><span class="p">):</span>
<span class="kn">import</span> <span class="nn">cv2</span>


<div class="viewcode-block" id="ImageFaceRatioFilter"><a class="viewcode-back" href="../../../../data_juicer.ops.filter.html#data_juicer.ops.filter.ImageFaceRatioFilter">[docs]</a><span class="nd">@OPERATORS</span><span class="o">.</span><span class="n">register_module</span><span class="p">(</span><span class="n">OP_NAME</span><span class="p">)</span>
<div class="viewcode-block" id="ImageFaceRatioFilter"><a class="viewcode-back" href="../../../../data_juicer.ops.filter.html#data_juicer.ops.filter.ImageFaceRatioFilter">[docs]</a><span class="nd">@UNFORKABLE</span><span class="o">.</span><span class="n">register_module</span><span class="p">(</span><span class="n">OP_NAME</span><span class="p">)</span>
<span class="nd">@OPERATORS</span><span class="o">.</span><span class="n">register_module</span><span class="p">(</span><span class="n">OP_NAME</span><span class="p">)</span>
<span class="nd">@LOADED_IMAGES</span><span class="o">.</span><span class="n">register_module</span><span class="p">(</span><span class="n">OP_NAME</span><span class="p">)</span>
<span class="k">class</span> <span class="nc">ImageFaceRatioFilter</span><span class="p">(</span><span class="n">Filter</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Filter to keep samples with face area ratios within a specific range.</span>
<span class="sd"> &quot;&quot;&quot;</span>

<span class="n">_default_kwargs</span> <span class="o">=</span> <span class="p">{</span><span class="s1">&#39;upsample_num_times&#39;</span><span class="p">:</span> <span class="mi">0</span><span class="p">}</span>
<span class="n">_default_kwargs</span> <span class="o">=</span> <span class="p">{</span>
<span class="s1">&#39;scaleFactor&#39;</span><span class="p">:</span> <span class="mf">1.1</span><span class="p">,</span>
<span class="s1">&#39;minNeighbors&#39;</span><span class="p">:</span> <span class="mi">3</span><span class="p">,</span>
<span class="s1">&#39;minSize&#39;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
<span class="s1">&#39;maxSize&#39;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
<span class="p">}</span>

<div class="viewcode-block" id="ImageFaceRatioFilter.__init__"><a class="viewcode-back" href="../../../../data_juicer.ops.filter.html#data_juicer.ops.filter.ImageFaceRatioFilter.__init__">[docs]</a> <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
<span class="n">cv_classifier</span><span class="o">=</span><span class="s1">&#39;&#39;</span><span class="p">,</span>
<span class="n">min_ratio</span><span class="p">:</span> <span class="n">ClosedUnitInterval</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
<span class="n">max_ratio</span><span class="p">:</span> <span class="n">ClosedUnitInterval</span> <span class="o">=</span> <span class="mf">0.4</span><span class="p">,</span>
<span class="n">any_or_all</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;any&#39;</span><span class="p">,</span>
Expand All @@ -117,6 +127,8 @@ <h1>Source code for data_juicer.ops.filter.image_face_ratio_filter</h1><div clas
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Initialization method.</span>

<span class="sd"> :param cv_classifier: OpenCV classifier path for face detection.</span>
<span class="sd"> By default, we will use &#39;haarcascade_frontalface_alt.xml&#39;.</span>
<span class="sd"> :param min_ratio: Min ratio for the largest face area in an image.</span>
<span class="sd"> :param max_ratio: Max ratio for the largest face area in an image.</span>
<span class="sd"> :param any_or_all: Keep this sample with &#39;any&#39; or &#39;all&#39; strategy of</span>
Expand All @@ -127,6 +139,10 @@ <h1>Source code for data_juicer.ops.filter.image_face_ratio_filter</h1><div clas
<span class="sd"> :param kwargs: Extra keyword arguments.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

<span class="k">if</span> <span class="n">cv_classifier</span> <span class="o">==</span> <span class="s1">&#39;&#39;</span><span class="p">:</span>
<span class="n">cv_classifier</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">cv2</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">haarcascades</span><span class="p">,</span>
<span class="s1">&#39;haarcascade_frontalface_alt.xml&#39;</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">min_ratio</span> <span class="o">=</span> <span class="n">min_ratio</span>
<span class="bp">self</span><span class="o">.</span><span class="n">max_ratio</span> <span class="o">=</span> <span class="n">max_ratio</span>

Expand All @@ -140,8 +156,8 @@ <h1>Source code for data_juicer.ops.filter.image_face_ratio_filter</h1><div clas
<span class="sa">f</span><span class="s1">&#39;Can only be one of [&quot;any&quot;, &quot;all&quot;].&#39;</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">any</span> <span class="o">=</span> <span class="p">(</span><span class="n">any_or_all</span> <span class="o">==</span> <span class="s1">&#39;any&#39;</span><span class="p">)</span>

<span class="c1"># Initialize face detector</span>
<span class="bp">self</span><span class="o">.</span><span class="n">detector</span> <span class="o">=</span> <span class="n">dlib</span><span class="o">.</span><span class="n">get_frontal_face_detector</span><span class="p">()</span></div>
<span class="bp">self</span><span class="o">.</span><span class="n">model_key</span> <span class="o">=</span> <span class="n">prepare_model</span><span class="p">(</span><span class="n">model_type</span><span class="o">=</span><span class="s1">&#39;opencv_classifier&#39;</span><span class="p">,</span>
<span class="n">model_path</span><span class="o">=</span><span class="n">cv_classifier</span><span class="p">)</span></div>

<div class="viewcode-block" id="ImageFaceRatioFilter.compute_stats"><a class="viewcode-back" href="../../../../data_juicer.ops.filter.html#data_juicer.ops.filter.ImageFaceRatioFilter.compute_stats">[docs]</a> <span class="k">def</span> <span class="nf">compute_stats</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">,</span> <span class="n">context</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
<span class="c1"># check if it&#39;s computed already</span>
Expand All @@ -159,25 +175,20 @@ <h1>Source code for data_juicer.ops.filter.image_face_ratio_filter</h1><div clas
<span class="n">sample</span><span class="p">,</span> <span class="n">images</span> <span class="o">=</span> <span class="n">load_data_with_context</span><span class="p">(</span><span class="n">sample</span><span class="p">,</span> <span class="n">context</span><span class="p">,</span>
<span class="n">loaded_image_keys</span><span class="p">,</span> <span class="n">load_image</span><span class="p">)</span>

<span class="n">model</span> <span class="o">=</span> <span class="n">get_model</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">model_key</span><span class="p">)</span>

<span class="c1"># detect faces</span>
<span class="n">face_detections</span> <span class="o">=</span> <span class="p">{}</span>
<span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">image</span> <span class="ow">in</span> <span class="n">images</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
<span class="n">img</span> <span class="o">=</span> <span class="n">pil_to_opencv</span><span class="p">(</span><span class="n">image</span><span class="p">)</span>
<span class="n">dets</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">detector</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="o">**</span><span class="bp">self</span><span class="o">.</span><span class="n">extra_kwargs</span><span class="p">)</span>
<span class="n">face_detections</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="p">[[</span>
<span class="nb">max</span><span class="p">(</span><span class="n">det</span><span class="o">.</span><span class="n">left</span><span class="p">(),</span> <span class="mi">0</span><span class="p">),</span>
<span class="nb">max</span><span class="p">(</span><span class="n">det</span><span class="o">.</span><span class="n">top</span><span class="p">(),</span> <span class="mi">0</span><span class="p">),</span>
<span class="nb">min</span><span class="p">(</span><span class="n">det</span><span class="o">.</span><span class="n">right</span><span class="p">(),</span> <span class="n">image</span><span class="o">.</span><span class="n">width</span><span class="p">),</span>
<span class="nb">min</span><span class="p">(</span><span class="n">det</span><span class="o">.</span><span class="n">bottom</span><span class="p">(),</span> <span class="n">image</span><span class="o">.</span><span class="n">height</span><span class="p">)</span>
<span class="p">]</span> <span class="k">for</span> <span class="n">det</span> <span class="ow">in</span> <span class="n">dets</span><span class="p">]</span>
<span class="n">face_detections</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">detect_faces</span><span class="p">(</span><span class="n">image</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span>
<span class="o">**</span><span class="bp">self</span><span class="o">.</span><span class="n">extra_kwargs</span><span class="p">)</span>
<span class="n">logger</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;detections: </span><span class="si">{</span><span class="n">face_detections</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>

<span class="c1"># compute face area ratios for each image considering the largest face</span>
<span class="n">face_area_ratios</span> <span class="o">=</span> <span class="p">{}</span>
<span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">dets</span> <span class="ow">in</span> <span class="n">face_detections</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
<span class="n">image_area</span> <span class="o">=</span> <span class="n">images</span><span class="p">[</span><span class="n">key</span><span class="p">]</span><span class="o">.</span><span class="n">width</span> <span class="o">*</span> <span class="n">images</span><span class="p">[</span><span class="n">key</span><span class="p">]</span><span class="o">.</span><span class="n">height</span>
<span class="n">face_area_ratios</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="nb">max</span><span class="p">([(</span><span class="n">x2</span> <span class="o">-</span> <span class="n">x1</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="n">y2</span> <span class="o">-</span> <span class="n">y1</span><span class="p">)</span>
<span class="k">for</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">y2</span> <span class="ow">in</span> <span class="n">dets</span><span class="p">],</span>
<span class="n">face_area_ratios</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="nb">max</span><span class="p">([</span><span class="n">w</span> <span class="o">*</span> <span class="n">h</span> <span class="k">for</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">h</span> <span class="ow">in</span> <span class="n">dets</span><span class="p">],</span>
<span class="n">default</span><span class="o">=</span><span class="mf">0.0</span><span class="p">)</span> <span class="o">/</span> <span class="n">image_area</span>
<span class="n">logger</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;ratios: </span><span class="si">{</span><span class="n">face_area_ratios</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>

Expand Down
Loading

0 comments on commit fec586d

Please sign in to comment.